[sparse] add trivial sparse rule for lax.copy_p

2025-04-19 05:16:06 +00:00 · 2022-08-19 14:25:41 -07:00 · 2022-08-19 14:25:41 -07:00 · d78b96937c
commit d78b96937c
parent bfef22ad7c
1 changed files with 1 additions and 0 deletions
--- a/jax/experimental/sparse/transform.py
+++ b/jax/experimental/sparse/transform.py
@ -78,6 +78,7 @@ _zero_preserving_unary_primitives = [
  lax.atan_p,
  lax.atanh_p,
  lax.bessel_i1e_p,
+  lax.copy_p,
  lax.expm1_p,
  lax.log1p_p,
  lax.neg_p,