Repository: google/jax

Issue:
jit changes behaviour of sparse.coo_matmat with transposed argument
JIT gives erroneous results for the multiplication of a sparse matrix with a transposed dense matrix.

```python

import jax
import jax.numpy as jnp
import jax.experimental.sparse as jsparse

d = jnp.r_[0.1, 2., 3., -2]
x = jax.random.normal(jax.random.PRNGKey(0), (3, 3))


def ok(d, x):
    i = jnp.r_[0, 0, 1, 2]
    j = jnp.r_[0, 2, 0, 0]
    shape = (3, 3)
    return jsparse.coo_matmat(d, i, j, x, shape=shape)


def NOT_ok(d, x):
    i = jnp.r_[0, 0, 1, 2]
    j = jnp.r_[0, 2, 0, 0]
    shape = (3, 3)
    return jsparse.coo_matmat(d, i, j, x.T, shape=shape)

e_ok = ok(d, x) - jax.jit(ok)(d, x)
e_not_ok = NOT_ok(d, x) - jax.jit(NOT_ok)(d, x)

print('ok: error = \n', e_ok)
print('\nNOT ok: error = \n',e_not_ok)
## OUTPUT
#ok: error = 
# [[0. 0. 0.]
# [0. 0. 0.]
# [0. 0. 0.]]
#
#NOT ok: error = 
# [[ 0.97765154  0.7773348  -0.0488826 ]
# [ 0.         -3.0031524  -1.4664774 ]
# [ 0.          2.0021017   0.9776516 ]]
```


Discussion:
Are you seeing this only on a GPU backend?
I don't reproduce the issue on CPU; I suspect it has to do with the GPU-only lowering to cusparse. We need to change the layout definition here for a transposed matrix: https://github.com/google/jax/blob/c94f41290cc5c28fba059e172279994107c3103e/jaxlib/cusparse.py#L285

But I'm not certain how to tell at this level whether the input buffer is transposed... perhaps XLA operators somehow expose their strides?
Yes, I'm seeing this on a GPU backend only.
I took a closer look at this and realized that in the transposed version, you're computing the matrix product of a (3,3) matrix and a (4, 3) matrix, which should be an error! The issue is poor input validation, which should be fixed by #7557.

The reason that there are different answers between GPU/JIT and non-jit approaches is that the two proceed based on different sets assumptions about the input, both of which are incorrect in their own way. With #7557, this should correctly result in an assertion error in both cases.
> I took a closer look at this and realized that in the transposed version, you're computing the matrix product of a (3,3) matrix and a (4, 3) matrix, which should be an error! The issue is poor input validation, which should be fixed by #7557.

Sorry, I think I was distracted while making this minimum working example, in my application where this bug first surfaced the dimensions matched and the JIT messed up the results. I've changed the dimension of the `x` matrix in the example to a square matrix and the error persists.

> The reason that there are different answers between GPU/JIT and non-jit approaches is that the two proceed based on different sets assumptions about the input, both of which are incorrect in their own way. With #7557, this should correctly result in an assertion error in both cases.

In the up side, I guess this helped improve the package in other ways with the validation.

Great, thanks for checking! I'll try to get to the bottom of the new & improved repro.
I think I found the issue – should be fixed by #7560

Resolving patch:
diff --git a/jaxlib/cusparse.py b/jaxlib/cusparse.py
--- a/jaxlib/cusparse.py
+++ b/jaxlib/cusparse.py
@@ -137,8 +137,9 @@ def csr_matmat(c, data, indices, indptr, B, *, shape, transpose=False, compute_d
   dtype = np.dtype(c.get_shape(data).element_type())
   index_dtype = np.dtype(c.get_shape(indices).element_type())
   B_dtype = np.dtype(c.get_shape(B).element_type())
+  B_shape = c.get_shape(B).dimensions()
   rows, cols = shape
-  _, Ccols = c.get_shape(B).dimensions()
+  _, Ccols = B_shape
   nnz, = c.get_shape(data).dimensions()
 
   if compute_dtype is None:
@@ -154,11 +155,10 @@ def csr_matmat(c, data, indices, indptr, B, *, shape, transpose=False, compute_d
       b"cusparse_csr_matmat",
       operands=(data, indices, indptr, B),
       operand_shapes_with_layout=(
-          # All are 1D, so no layout necessary
           c.get_shape(data),
           c.get_shape(indices),
           c.get_shape(indptr),
-          c.get_shape(B),
+          _Shape.array_shape(B_dtype, B_shape, (1, 0)),
       ),
       shape_with_layout=_Shape.tuple_shape((
           _Shape.array_shape(compute_dtype, (out_size, Ccols), (1, 0)),
@@ -272,8 +272,9 @@ def coo_matmat(c, data, row, col, B, *, shape, transpose=False, compute_dtype=No
   dtype = np.dtype(c.get_shape(data).element_type())
   index_dtype = np.dtype(c.get_shape(row).element_type())
   B_dtype = np.dtype(c.get_shape(B).element_type())
+  B_shape = c.get_shape(B).dimensions()
   rows, cols = shape
-  _, Ccols = c.get_shape(B).dimensions()
+  _, Ccols = B_shape
   nnz, = c.get_shape(data).dimensions()
 
   if compute_dtype is None:
@@ -289,11 +290,10 @@ def coo_matmat(c, data, row, col, B, *, shape, transpose=False, compute_dtype=No
       b"cusparse_coo_matmat",
       operands=(data, row, col, B),
       operand_shapes_with_layout=(
-          # All are 1D, so no layout necessary
           c.get_shape(data),
           c.get_shape(row),
           c.get_shape(col),
-          c.get_shape(B),
+          _Shape.array_shape(B_dtype, B_shape, (1, 0)),
       ),
       shape_with_layout=_Shape.tuple_shape((
           _Shape.array_shape(compute_dtype, (out_size, Ccols), (1, 0)),


Test patch:
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -202,7 +202,6 @@ def test_coo_matvec(self, shape, dtype, transpose):
     self.assertAllClose(op(M) @ v, matvec(*args), rtol=MATMUL_TOL)
     self.assertAllClose(op(M) @ v, jit(matvec)(*args), rtol=MATMUL_TOL)
 
-  @unittest.skipIf(jtu.device_under_test() != "gpu", "test requires GPU")
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_T={}".format(jtu.format_shape_dtype_string(shape, dtype), transpose),
        "shape": shape, "dtype": dtype, "transpose": transpose}
@@ -229,6 +228,23 @@ def test_coo_matmat(self, shape, dtype, transpose):
     y, dy = jvp(lambda x: sparse.coo_matmat(x, M.row, M.col, B, shape=shape, transpose=transpose).sum(), (M.data, ), (jnp.ones_like(M.data), ))
     self.assertAllClose((op(M) @ B).sum(), y, rtol=MATMUL_TOL)
 
+  def test_coo_matmat_layout(self):
+    # Regression test for https://github.com/google/jax/issues/7533
+    d = jnp.array([1.0, 2.0, 3.0, 4.0])
+    i = jnp.array([0, 0, 1, 2])
+    j = jnp.array([0, 2, 0, 0])
+    shape = (3, 3)
+
+    x = jnp.arange(9).reshape(3, 3).astype(d.dtype)
+
+    def f(x):
+      return sparse.coo_matmat(d, i, j, x.T, shape=shape)
+
+    result = f(x)
+    result_jit = jit(f)(x)
+
+    self.assertAllClose(result, result_jit)
+
   @unittest.skipIf(jtu.device_under_test() != "gpu", "test requires GPU")
   def test_gpu_translation_rule(self):
     version = xla_bridge.get_backend().platform_version