Repository: google/jax
Issue: [sparse] Difference between BCOO matrices removes zeros

### Description

The difference operation between two `BCOO` sparse matrices implicitly removes zeros. The behaviour is different from the one observed for the addition. Consider the example below.

```python
import jax.experimental.sparse as jsp
import jax.numpy as jnp

a = jsp.BCOO((jnp.array([1.0]), jnp.array([[0, 0]])), shape=(3, 3))
b_values = jnp.array([0.0, 1.0, 0.0])
b_indices = jnp.array([[0, 0], [1, 1], [2, 2]])
b = jsp.BCOO((b_values, b_indices), shape=(3, 3))

print("Difference\\n", (a - b).indices)
print("Addition\\n", (a + b).indices)
```

The output is 
```
Difference
 [[0 0]
 [1 1]
 [3 3]
 [3 3]]
Addition
 [[0 0]
 [0 0]
 [1 1]
 [2 2]]
```

As you can see, in the difference operation the indices associated to null values in `a` or `b` have disappeared and they have been replaced with the out-of-bounds index `[3,3]`.

The expected behaviour can be obtained using `a + (-1)*b`.

### What jax/jaxlib version are you using?

jax v0.3.25, jaxlib v.0.3.25

### Which accelerator(s) are you using?

CPU/GPU

### Additional system info

Python 3.10.8, ArchLinux

### NVIDIA GPU info

NVIDIA-SMI 520.56.06    Driver Version: 520.56.06    CUDA Version: 11.8 

Changed files:
tests/sparsify_test.py, jax/experimental/sparse/transform.py

diff --git a/jax/experimental/sparse/transform.py b/jax/experimental/sparse/transform.py
index b731d9426..f593e4174 100644
--- a/jax/experimental/sparse/transform.py
+++ b/jax/experimental/sparse/transform.py
@@ -70,6 +70,13 @@ from jax.experimental.sparse import BCOO
 
 sparse_rules : Dict[core.Primitive, Callable] = {}
 
+_zero_preserving_linear_unary_primitives = [
+  lax.copy_p,
+  lax.imag_p,
+  lax.neg_p,
+  lax.real_p,
+]
+
 _zero_preserving_unary_primitives = [
   lax.abs_p,
   lax.asin_p,
@@ -77,19 +84,15 @@ _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,
-  lax.real_p,
-  lax.imag_p,
   lax.sign_p,
   lax.sin_p,
   lax.sinh_p,
   lax.sqrt_p,
   lax.tan_p,
   lax.tanh_p,
-  lax.convert_element_type_p
+  lax.convert_element_type_p,
 ]
 
 _densifying_primitives : List[core.Primitive] = [
@@ -462,15 +465,17 @@ def _ensure_unique_indices(spenv, spvalue):
   if spvalue.is_dense() or spvalue.unique_indices:
     return spvalue
   arr = spvalues_to_arrays(spenv, spvalue)
-  arr = arr.sum_duplicates(nse=arr.nse)
+  arr = arr.sum_duplicates(nse=arr.nse, remove_zeros=False)
   return arrays_to_spvalues(spenv, arr)
 
-def _zero_preserving_unary_op(prim):
+def _zero_preserving_unary_op(prim, linear):
   def func(spenv, *spvalues, **kwargs):
     assert len(spvalues) == 1
-    # Since unary operations don't commute with addition, we need to ensure
-    # that indices are unique before applying the operator elementwise.
-    spvalue = _ensure_unique_indices(spenv, spvalues[0])
+    spvalue = spvalues[0]
+    if not linear:
+      # For non-linear unary operations, we need to ensure that
+      # indices are unique before applying the operator elementwise.
+      spvalue = _ensure_unique_indices(spenv, spvalue)
     buf = spenv.data(spvalue)
     buf_out = prim.bind(buf, **kwargs)
     if spvalues[0].is_sparse():
@@ -484,7 +489,9 @@ def _zero_preserving_unary_op(prim):
   return func
 
 for _prim in _zero_preserving_unary_primitives:
-  sparse_rules[_prim] = _zero_preserving_unary_op(_prim)
+  sparse_rules[_prim] = _zero_preserving_unary_op(_prim, linear=False)
+for _prim in _zero_preserving_linear_unary_primitives:
+  sparse_rules[_prim] = _zero_preserving_unary_op(_prim, linear=True)
 
 def _standard_sparse_rule(prim, sparse_op):
   def _sparse_rule(spenv, *spvalues, **kwds):
diff --git a/tests/sparsify_test.py b/tests/sparsify_test.py
index f9a0a331f..f3576b9a9 100644
--- a/tests/sparsify_test.py
+++ b/tests/sparsify_test.py
@@ -500,6 +500,7 @@ class SparsifyTest(jtu.JaxTestCase):
   @parameterized.named_parameters(
       {"testcase_name": f"_{op.__name__}", "op": op, "dtype": dtype, "kwds": kwds}
       for op, dtype, kwds in [
+        (jnp.copy, jnp.float32, {}),
         (lax.abs, jnp.float32, {}),
         (lax.asin, jnp.float32, {}),
         (lax.asinh, jnp.float32, {}),
@@ -530,10 +531,14 @@ class SparsifyTest(jtu.JaxTestCase):
     indices = rng_idx((nse, len(shape)), jnp.int32)
     mat = BCOO((data, indices), shape=shape)
 
-    sparse_result = self.sparsify(partial(op, **kwds))(mat).todense()
+    sparse_result = self.sparsify(partial(op, **kwds))(mat)
     dense_result = op(mat.todense(), **kwds)
 
-    self.assertArraysAllClose(sparse_result, dense_result)
+    self.assertArraysAllClose(sparse_result.todense(), dense_result)
+
+    # Ops that commute with addition should not deduplicate indices.
+    if op in [jnp.copy, lax.neg, lax.real, lax.imag]:
+      self.assertArraysAllClose(sparse_result.indices, indices)
 
 
 class SparsifyTracerTest(SparsifyTest):