Merge pull request #14500 from jakevdp:bcsr-matmul-test

PiperOrigin-RevId: 510034750
jax-ml · Feb 16, 2023 · d8514d0 · d8514d0
2 parents 0af9fff + 29f91c5
commit d8514d0
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
@@ -2227,9 +2227,8 @@ def test_bcsr_dot_general(self, dtype: np.dtype, props: BatchedDotGeneralPropert
       # Dense dimensions not yet fully supported in reverse mode.
       modes = ['fwd'] if props.n_dense != 0 else ['fwd', 'rev']
       self._CheckGradsSparse(dense_fun, sparse_fun, args_maker, modes=modes, atol=tol, rtol=tol)
-    # TODO: add this once bcsr_broadcast_in_dim & bcsr_concatenate are implemented
-    # self._CheckBatchingSparse(dense_fun, sparse_fun, args_maker, atol=tol, rtol=tol,
-    #                           bdims=self._random_bdims(props.n_batch, len(props.rhs_shape)))
+    self._CheckBatchingSparse(dense_fun, sparse_fun, args_maker, atol=tol, rtol=tol,
+                              bdims=self._random_bdims(props.n_batch, len(props.rhs_shape)))
 
   @jtu.sample_product(
     [dict(shape=shape, n_batch=layout.n_batch, n_dense=layout.n_dense)