Repository: google/jax

Issue:
[sparse] Wrong jacobian computation for bcoo_spdot_general
It seems that the Jacobian computed by `bcoo_spdot_general`, invoked through the operator `@` does not organize the values in the correct order.

In order to show the bug, I have written the following code.
```python
import jax
import jax.numpy as jnp
from jax.experimental.sparse.bcoo import BCOO

A_indices = jnp.array([[0, 1], [0, 2], [1, 1], [1, 2], [1, 0]])
A_values = jnp.array([-2.0, 1.0, -1.0, 0.5, 2.0])
A_shape = (2, 3)

B_indices = jnp.array([[0, 2], [2, 1], [0, 3], [1, 3], [1, 0], [0, 0]])
B_values = jnp.array([10.0, 100.0, 1000.0, -5.0, -50.0, -500.0])
B_shape = (3, 4)

def sp_sp_product(v1, v2):
    sparse1 = BCOO((v1, A_indices), shape=A_shape)
    sparse2 = BCOO((v2, B_indices), shape=B_shape)
    return (sparse1 @ sparse2).todense()

def sp_dense_product(v1, v2):
    sparse1 = BCOO((v1, A_indices), shape=A_shape)
    dense2 = BCOO((v2, B_indices), shape=B_shape).todense()
    return sparse1 @ dense2

sp_sp_product_jac = jax.jacfwd(sp_sp_product, argnums=1)(A_values, B_values)
sparse_dense_product_jac = jax.jacfwd(sp_dense_product, argnums=1)(A_values, B_values)

assert jnp.array_equal(
    sp_sp_product_jac, sparse_dense_product_jac
), "The two Jacobian should be the same"
```

`sparse_dense_product_jac` is correct (as it can be verified analitically), while `sp_sp_product_jac` is incorrect since it should be identical to `sparse_dense_product_jac` but it is not.
The shape of `sp_sp_product_jac` in the expected one and the values in `sp_sp_product_jac` seems correct but they are not placed in the expected location. Probably there is a problem in the indices of the sparse jacobian matrix.



Discussion:
@jakevdp As the author of https://github.com/google/jax/pull/8848, have you any thought about it?
I'll take a look
Still working on debugging this... it's a tricky one.
Inspecting the code it seems that the problem is that the indices relative to the output JVP are assumed to be the indices of the output of the result of the product `bcoo_spdot_general`.

I have tried to return these indices modifying the function like that.

```python
def _bcoo_spdot_general_jvp(primals, tangents, **kwds):
    lhs_data, lhs_indices, rhs_data, rhs_indices = primals
    lhs_data_dot, lhs_indices_dot, rhs_data_dot, rhs_indices_dot = tangents
    primals_out = _bcoo_spdot_general(*primals, **kwds)
    assert type(lhs_indices_dot) is ad.Zero
    assert type(rhs_indices_dot) is ad.Zero
    if type(lhs_data_dot) is not ad.Zero:
        data_dot_out, indices_dot_out = _bcoo_spdot_general(
            lhs_data_dot, lhs_indices, rhs_data, rhs_indices, **kwds
        )
    if type(rhs_data_dot) is not ad.Zero:
        data_dot_out, indices_dot_out = _bcoo_spdot_general(
            lhs_data, lhs_indices, rhs_data_dot, rhs_indices, **kwds
        )
    return primals_out, [data_dot_out, indices_dot_out]
```

This code does not crash, but the tangent indices `indices_dot_out` are not passed to the JVP function of `bcoo_todense` (that makes the same assumption that tangent and primal data vectors share the same indices)
https://github.com/google/jax/blob/3184dd65a222354bffa2466d9a375162f5649132/jax/experimental/sparse/bcoo.py#L256
so the final result is identical.
The assumption that the tangent and primal vectors share the same indices is fundamental to how autodiff of sparse objects works (I'm not sure how you could have nonzero gradients in non-existent values!). I'm not sure that your modified function is the correct approach, because the gradient of integer indices is not those indices themselves, but must be a symbolic zero.

My suspicion at the moment is that this has something to do with the interaction between batching and autodiff that happens within a jacobian... we end up with a sparse array with batched data and unbatched indices, and I think our test coverage is not all that extensive for that case.
OK, I think the issue comes down to this line in the spdot_general implementation: https://github.com/google/jax/blob/5522ed1702de3e806ecb41b7adc00f9841e27a62/jax/experimental/sparse/bcoo.py#L1090

`sum_duplicates` sorts the indices and moves explicit zero-values to the end - since explicit zeros are different in each batch, it has to return a different set of indices for each batch, such that the output data no longer shares indices. We need a version of `sum_duplicates` that works correctly when data is batched and indices are not.
Put another way, this issue has brought up a more general concern that I hadn't considered before: sparse autodiff is fundamentally incompatible with any operation in which the output indices depend on the values in the data, because that may break the property that primals and tangents must share indices.
Here's a shorter example showing how grad of `sum_duplicates` is incorrect if there are zeros in the data. I believe this is another effect of the root cause of this bug:
```python
from functools import partial
import jax
import jax.numpy as jnp
from jax.experimental import sparse

indices = jnp.array([[1], [2]])
data = jnp.array([0.0, 1.0])

def f(data, sum_duplicates = True):
  mat = sparse.BCOO((data, indices), shape=(3,))
  if sum_duplicates:
    mat = mat.sum_duplicates()
  return mat.todense().sum()

print(jax.grad(partial(f, sum_duplicates=True))(data))
# [0. 1.]
print(jax.grad(partial(f, sum_duplicates=False))(data))
# [1. 1.]
```
It seems to me that you changed your mind with respect to the cause in the last comment, but I reply to your previous comment sharing my current knowledge.

> My suspicion at the moment is that this has something to do with the interaction between batching and autodiff that happens within a jacobian... we end up with a sparse array with batched data and unbatched indices, and I think our test coverage is not all that extensive for that case.

The error can be reproduced for the JVP only, without the batching involved in the computation of the Jacobian.

```python
import jax
import jax.numpy as jnp
from jax.experimental.sparse.bcoo import BCOO

A_indices = jnp.array([[0, 1], [0, 2], [1, 1], [1, 2], [1, 0]])
A_values = jnp.array([-2.0, 1.0, -1.0, 0.5, 2.0])
A_shape = (2, 3)

B_indices = jnp.array([[0, 2], [2, 1], [0, 3], [1, 3], [1, 0], [0, 0]])
B_values = jnp.array([10.0, 100.0, 1000.0, -5.0, -50.0, -500.0])
B_shape = (3, 4)


def sp_sp_product(b_v):
    sparse1 = BCOO((A_values, A_indices), shape=A_shape)
    sparse2 = BCOO((b_v, B_indices), shape=B_shape)
    return (sparse1 @ sparse2).todense()


def sp_dense_product(b_v):
    sparse1 = BCOO((A_values, A_indices), shape=A_shape)
    dense2 = BCOO((b_v, B_indices), shape=B_shape).todense()
    return sparse1 @ dense2


first_row_vector = jnp.eye(B_values.size)[0]

_, sp_sp_product_jvp = jax.jvp(sp_sp_product, (B_values,), (first_row_vector,))
_, sp_dense_product_jvp = jax.jvp(sp_dense_product, (B_values,), (first_row_vector,))

assert jnp.array_equal(
    sp_sp_product_jvp, sp_dense_product_jvp
), "The two JVP should be the same"
```
> Put another way, this issue has brought up a more general concern that I hadn't considered before: sparse autodiff is fundamentally incompatible with any operation in which the output indices depend on the values in the data, because that may break the property that primals and tangents must share indices.

I totally agree with this!
In my opinion, a workaround is to remove the constraint that the tangent vector of indices is Zero and use it to store the indices. I know that it is not mathematically correct, but if this choice stays hidden in the code maybe it causes no harm.

I'm looking into it to see if that could work in practice without a major refactor of the API.
> I'm not sure that your modified function is the correct approach, because the gradient of integer indices is not those indices themselves, but must be a symbolic zero.

You are right, this constraint is hard to circumvent because it is deeply integrated into JAX.
I have written a version of the code that does not causes the bug for the examples I posted, but the code breaks in other tests linked to sparse matrices because when the tangent of the indices isn't zero, then an assertion fails in
https://github.com/google/jax/blob/0c02f7935aa90d1d895f34882186e889e2b8923e/jax/interpreters/ad.py#L454

I'm willing to help fix the problem, but I do not have sufficient knowledge of JAX to come up with a solution that is consistent with the established inner workings of the library.
I don't think this will require any deep fix, we just need to not use a problematic version of `sum_duplicates` in the `bcoo_spdot_general` implementation. I think the solution will be to split the current `sum_duplicates` utility into a couple utilities, or perhaps one with several options: `sum_duplicates`, `sort_indices`, and `eliminate_zeros`. The first two are fine for autodiff, the third should error on autodiff. Then in functions that we want to be differentiable, we don't call the third.

Resolving patch:
diff --git a/jax/experimental/sparse/bcoo.py b/jax/experimental/sparse/bcoo.py
--- a/jax/experimental/sparse/bcoo.py
+++ b/jax/experimental/sparse/bcoo.py
@@ -84,6 +84,8 @@ def _bcoo_nse(mat, n_batch=0, n_dense=0):
   mask = mask.sum(list(range(n_batch, mask.ndim)))
   return mask.max()
 
+# TODO(jakevdp): add a custom autodiff rule that errors if remove_zeros=True, because
+# it produces wrong values. See https://github.com/google/jax/issues/10163
 def _bcoo_sum_duplicates(data, indices, shape, nse=None, remove_zeros=True):
   if nse is None and isinstance(jnp.array(0), core.Tracer):
     raise ValueError("When used with JIT, vmap, or another transform, sum_duplicates() "
@@ -1229,7 +1231,8 @@ def _bcoo_spdot_general_unbatched(lhs_data, lhs_indices, rhs_data, rhs_indices,
   out_indices = out_indices.at[:, :, lhs_j.shape[-1]:].set(rhs_j[None, :])
   out_indices = out_indices.reshape(len(out_data), out_indices.shape[-1])
   out_nse = (lhs.nse if lhs_j.shape[1] else 1) * (rhs.nse if rhs_j.shape[1] else 1)
-  return _bcoo_sum_duplicates(out_data, out_indices, out_shape, nse=out_nse)
+  # Note: remove_zeros=True is incompatible with autodiff.
+  return _bcoo_sum_duplicates(out_data, out_indices, out_shape, nse=out_nse, remove_zeros=False)
 
 @bcoo_spdot_general_p.def_impl
 def _bcoo_spdot_general_impl(lhs_data, lhs_indices, rhs_data, rhs_indices, *, lhs_spinfo: BCOOInfo, rhs_spinfo: BCOOInfo, dimension_numbers):
@@ -1770,7 +1773,8 @@ def sum_duplicates(self, nse=None, remove_zeros=True):
         If it is smaller than the number required, data will be silently discarded.
       remove_zeros : bool (default=True). If True, remove explicit zeros from the data
         as part of summing duplicates. If False, then explicit zeros at unique indices
-        will remain among the specified elements.
+        will remain among the specified elements. Note: remove_zeros=True is incompatible
+        with autodiff.
     """
     data, indices = _bcoo_sum_duplicates(self.data, self.indices, self.shape,
                                          nse=nse, remove_zeros=remove_zeros)


Test patch:
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -1454,6 +1454,38 @@ def f_sparse(lhs_data, rhs_data):
     self.assertAllClose(jf_dense_0, jf_sparse_0, rtol=tol)
     self.assertAllClose(jf_dense_1, jf_sparse_1, rtol=tol)
 
+  def test_bcoo_spdot_general_ad_bug(self):
+    # Regression test for https://github.com/google/jax/issues/10163
+    A_indices = jnp.array([[0, 1], [0, 2], [1, 1], [1, 2], [1, 0]])
+    A_values = jnp.array([-2.0, 1.0, -1.0, 0.5, 2.0])
+    A_shape = (2, 3)
+
+    B_indices = jnp.array([[0, 2], [2, 1], [0, 3], [1, 3], [1, 0], [0, 0]])
+    B_values = jnp.array([10.0, 100.0, 1000.0, -5.0, -50.0, -500.0])
+    B_shape = (3, 4)
+
+    def sp_sp_product(v1, v2):
+        A = sparse.BCOO((v1, A_indices), shape=A_shape)
+        B = sparse.BCOO((v2, B_indices), shape=B_shape)
+        return (A @ B).todense()
+
+    def sp_de_product(v1, v2):
+        A = sparse.BCOO((v1, A_indices), shape=A_shape)
+        B = sparse.BCOO((v2, B_indices), shape=B_shape).todense()
+        return A @ B
+
+    def de_de_product(v1, v2):
+        sparse1 = sparse.BCOO((v1, A_indices), shape=A_shape).todense()
+        dense2 = sparse.BCOO((v2, B_indices), shape=B_shape).todense()
+        return sparse1 @ dense2
+
+    sp_sp_jac = jax.jacfwd(sp_sp_product, argnums=1)(A_values, B_values)
+    sp_de_jac = jax.jacfwd(sp_de_product, argnums=1)(A_values, B_values)
+    de_de_jac = jax.jacfwd(de_de_product, argnums=1)(A_values, B_values)
+
+    self.assertAllClose(sp_sp_jac, de_de_jac)
+    self.assertAllClose(sp_de_jac, de_de_jac)
+
   @unittest.skipIf(jtu.device_under_test() == "tpu", "TPU has insufficient precision")
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}[n_batch={}]_{}[n_batch={}]_in_axes={}".format(