[sparse] support type promotion in CSR/COO matmul

jax/experimental/sparse/ops.py

@@ -45,6 +45,7 @@
 from jax._src.lib import cusparse
 from jax._src.lib import xla_bridge
 from jax._src.lib import xla_client
+from jax._src.numpy.lax_numpy import _promote_dtypes
 import jax.numpy as jnp
 
 xb = xla_bridge
@@ -807,10 +808,12 @@ def todense(self):
     return csr_todense(self.data, self.indices, self.indptr, shape=self.shape)
 
   def matvec(self, v):
-    return csr_matvec(self.data, self.indices, self.indptr, v, shape=self.shape)
+    data, v = _promote_dtypes(self.data, v)
+    return csr_matvec(data, self.indices, self.indptr, v, shape=self.shape)
 
   def matmat(self, B):
-    return csr_matmat(self.data, self.indices, self.indptr, B, shape=self.shape)
+    data, B = _promote_dtypes(self.data, B)
+    return csr_matmat(data, self.indices, self.indptr, B, shape=self.shape)
 
   def transpose(self, axes=None):
     assert axes is None
@@ -844,10 +847,12 @@ def todense(self):
     return csr_todense(self.data, self.indices, self.indptr, shape=self.shape[::-1]).T
 
   def matvec(self, v):
-    return csr_matvec(self.data, self.indices, self.indptr, v, shape=self.shape[::-1], transpose=True)
+    data, v = _promote_dtypes(self.data, v)
+    return csr_matvec(data, self.indices, self.indptr, v, shape=self.shape[::-1], transpose=True)
 
   def matmat(self, B):
-    return csr_matmat(self.data, self.indices, self.indptr, B, shape=self.shape[::-1], transpose=True)
+    data, B = _promote_dtypes(self.data, B)
+    return csr_matmat(data, self.indices, self.indptr, B, shape=self.shape[::-1], transpose=True)
 
   def transpose(self, axes=None):
     assert axes is None
@@ -881,10 +886,12 @@ def todense(self):
     return coo_todense(self.data, self.row, self.col, shape=self.shape)
 
   def matvec(self, v):
-    return coo_matvec(self.data, self.row, self.col, v, shape=self.shape)
+    data, v = _promote_dtypes(self.data, v)
+    return coo_matvec(data, self.row, self.col, v, shape=self.shape)
 
   def matmat(self, B):
-    return coo_matmat(self.data, self.row, self.col, B, shape=self.shape)
+    data, B = _promote_dtypes(self.data, B)
+    return coo_matmat(data, self.row, self.col, B, shape=self.shape)
 
   def transpose(self, axes=None):
     assert axes is None

tests/sparse_test.py

@@ -1488,9 +1488,15 @@ def test_matmul(self, shape, dtype, Obj, bshape):
     rng_b = jtu.rand_default(self.rng())
     M = rng(shape, dtype)
     Msp = Obj.fromdense(M)
+
+    # Test matching type
     x = rng_b(bshape, dtype)
     x = jnp.asarray(x)
+    self.assertAllClose(M @ x, Msp @ x, rtol=MATMUL_TOL)
 
+    # Test mismatched type
+    x = rng_b(bshape, np.int32)
+    x = jnp.asarray(x)
     self.assertAllClose(M @ x, Msp @ x, rtol=MATMUL_TOL)
 
   @parameterized.named_parameters(jtu.cases_from_list(