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 


Discussion:
Hi - thanks for the report. The issue isn't in subtraction, but rather the sparse rule for `neg_p`. Currently it's treated as any other unary operation, and in general unary operations do not commute with addition (i.e. `op(x + y) != op(x) + op(y)` in general for operations like `sin` and `cos`). For this reason, the unary operation sparse rule will deduplicate indices, which has the side effect of (1) sorting the indices, and (2) padding-out explicit zeros. This is fine in general if you consider the sparse data structure only as an abstract representation of a dense matrix.

Is it important in your application for explicit zeros to be maintained during operations? The sparse design currently does not include that in its design considerations, and so I expect you'll run into this kind of thing frequently if you depend on details of the sparse representation.

So, that said, I think this is "working as intended" for now, but we may be able to think about change the design goals of the sparse functionality if it's required for important use-cases. For the case of `neg_p` in particular, I think it would make sense to avoid re-indexing because it does commute with addition. What do you think?
Thank you for the clarification on the reasons leading to this behaviour!
In my application, I use BCOO matrices as "dictionaries" where the listed indices can be associated even to a null value, while non-listed indices are considered as undefined. I believe that this use case can be quite common: for example if we want to represent scalar values associated to points scattered on a N-dimensional grid, we can use a sparse matrix to represent points and values and we want to allow the null value.

Surely, I think that the `neg_p` behaviour should be consistent with the sum as you propose!
For the other operations, I agree that the duplicates should be removed, but I'm not sure that zeros should be removed too. In particular, for functions like `cos` (`cos(0) != 0`) that transform null values in non-null values, if the code does not explicitly prescribe to remove zeros, maybe the author expects that zero values will become non-null values after the transformation.



Makes sense - so I think we could call `sum_duplicates` with `remove_zeros=False` for unary ops, and also specialize `lax.neg_p` to avoid index modification.

Also, regarding cosine, I shouldn't have brought that up because unary ops that map zeros to non-zero values are not implemented in sparse (because the result is generally a dense matrix).
#13300 should address these issues - thanks for raising them!

Resolving patch:
diff --git a/jax/experimental/sparse/transform.py b/jax/experimental/sparse/transform.py
--- a/jax/experimental/sparse/transform.py
+++ b/jax/experimental/sparse/transform.py
@@ -70,6 +70,13 @@
 
 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 @@
   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 func(spenv, *spvalues, **kwargs):
   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):


Test patch:
diff --git a/tests/sparsify_test.py b/tests/sparsify_test.py
--- a/tests/sparsify_test.py
+++ b/tests/sparsify_test.py
@@ -500,6 +500,7 @@ def testWeakTypes(self):
   @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 @@ def testUnaryOperationsNonUniqueIndices(self, op, dtype, kwds):
     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):