Merge pull request #20673 from Micky774:sparse_test

PiperOrigin-RevId: 623326571

Author: jax authors

jax/experimental/sparse/test_util.py

@@ -41,10 +41,6 @@
     np.complex128: 1e-10,
 }
 
-GPU_LOWERING_ENABLED = gpu_sparse and (
-    gpu_sparse.cuda_is_supported or gpu_sparse.rocm_is_supported
-)
-
 
 def is_sparse(x):
   return isinstance(x, sparse.JAXSparse)

tests/sparse_bcoo_bcsr_test.py

@@ -151,7 +151,7 @@ def _is_required_cuda_version_satisfied(cuda_version):
 class BCOOTest(sptu.SparseTestCase):
 
   def gpu_matmul_warning_context(self, msg):
-    if sptu.GPU_LOWERING_ENABLED and config.jax_bcoo_cusparse_lowering:
+    if config.jax_bcoo_cusparse_lowering:
       return self.assertWarnsRegex(sparse.CuSparseEfficiencyWarning, msg)
     return contextlib.nullcontext()
 
@@ -479,9 +479,6 @@ def test_bcoo_dot_general(
       dtype=jtu.dtypes.floating + jtu.dtypes.complex,
   )
   @jax.default_matmul_precision("float32")
-  @unittest.skipIf(
-      not sptu.GPU_LOWERING_ENABLED, "test requires cusparse/hipsparse"
-  )
   @jtu.run_on_devices("gpu")
   def test_bcoo_dot_general_cusparse(
       self, lhs_shape, rhs_shape, dtype, lhs_contracting, rhs_contracting
@@ -528,9 +525,6 @@ def f_sparse(lhs_bcoo, lhs, rhs):
       dtype=jtu.dtypes.floating + jtu.dtypes.complex,
   )
   @jax.default_matmul_precision("float32")
-  @unittest.skipIf(
-      not sptu.GPU_LOWERING_ENABLED, "test requires cusparse/hipsparse"
-  )
   @jtu.run_on_devices("gpu")
   def test_bcoo_batched_matmat_cusparse(
       self,
@@ -581,9 +575,6 @@ def f_sparse(lhs_bcoo, lhs, rhs):
       ],
       dtype=jtu.dtypes.floating + jtu.dtypes.complex,
   )
-  @unittest.skipIf(
-      not sptu.GPU_LOWERING_ENABLED, "test requires cusparse/hipsparse"
-  )
   @jtu.run_on_devices("gpu")
   def test_bcoo_batched_matmat_default_lowering(
       self,
@@ -615,9 +606,6 @@ def test_bcoo_batched_matmat_default_lowering(
     matmat_default_lowering_fallback = sp_matmat(lhs_bcoo, rhs)
     self.assertArraysEqual(matmat_expected, matmat_default_lowering_fallback)
 
-  @unittest.skipIf(
-      not sptu.GPU_LOWERING_ENABLED, "test requires cusparse/hipsparse"
-  )
   @jtu.run_on_devices("gpu")
   def test_bcoo_dot_general_oob_and_unsorted_indices_cusparse(self):
     """Tests bcoo dot general with out-of-bound and unsorted indices."""

tests/sparse_test.py

@@ -16,7 +16,6 @@
 from functools import partial
 import itertools
 import math
-import unittest
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -351,9 +350,6 @@ def test_coo_sorted_indices(self):
     mat_resorted = mat_unsorted._sort_indices()
     self.assertArraysEqual(mat.todense(), mat_resorted.todense())
 
-  @unittest.skipIf(
-      not sptu.GPU_LOWERING_ENABLED, "test requires cusparse/hipsparse"
-  )
   @jtu.run_on_devices("gpu")
   def test_coo_sorted_indices_gpu_lowerings(self):
     dtype = jnp.float32
@@ -544,9 +540,7 @@ def test_coo_matmul_ad(self, shape, dtype, bshape):
       dtype=_lowerings.SUPPORTED_DATA_DTYPES,
       transpose=[True, False],
   )
-  @unittest.skipIf(
-      not sptu.GPU_LOWERING_ENABLED, "test requires cusparse/hipsparse"
-  )
+  @jtu.run_on_devices("gpu")
   def test_coo_spmv(self, shape, dtype, transpose):
     rng_sparse = sptu.rand_sparse(self.rng())
     rng_dense = jtu.rand_default(self.rng())
@@ -569,9 +563,7 @@ def test_coo_spmv(self, shape, dtype, transpose):
       dtype=_lowerings.SUPPORTED_DATA_DTYPES,
       transpose=[True, False],
   )
-  @unittest.skipIf(
-      not sptu.GPU_LOWERING_ENABLED, "test requires cusparse/hipsparse"
-  )
+  @jtu.run_on_devices("gpu")
   def test_coo_spmm(self, shape, dtype, transpose):
     rng_sparse = sptu.rand_sparse(self.rng())
     rng_dense = jtu.rand_default(self.rng())
@@ -594,9 +586,7 @@ def test_coo_spmm(self, shape, dtype, transpose):
       dtype=_lowerings.SUPPORTED_DATA_DTYPES,
       transpose=[True, False],
   )
-  @unittest.skipIf(
-      not sptu.GPU_LOWERING_ENABLED, "test requires cusparse/hipsparse"
-  )
+  @jtu.run_on_devices("gpu")
   def test_csr_spmv(self, shape, dtype, transpose):
     rng_sparse = sptu.rand_sparse(self.rng())
     rng_dense = jtu.rand_default(self.rng())
@@ -617,9 +607,7 @@ def test_csr_spmv(self, shape, dtype, transpose):
       dtype=_lowerings.SUPPORTED_DATA_DTYPES,
       transpose=[True, False],
   )
-  @unittest.skipIf(
-      not sptu.GPU_LOWERING_ENABLED, "test requires cusparse/hipsparse"
-  )
+  @jtu.run_on_devices("gpu")
   def test_csr_spmm(self, shape, dtype, transpose):
     rng_sparse = sptu.rand_sparse(self.rng())
     rng_dense = jtu.rand_default(self.rng())
@@ -1083,8 +1071,6 @@ class SparseSolverTest(sptu.SparseTestCase):
   )
   @jtu.run_on_devices("cpu", "cuda")
   def test_sparse_qr_linear_solver(self, size, reorder, dtype):
-    if jtu.test_device_matches(["cuda"]) and not sptu.GPU_LOWERING_ENABLED:
-      raise unittest.SkipTest('test requires cusparse/cusolver')
     rng = sptu.rand_sparse(self.rng())
     a = rng((size, size), dtype)
     nse = (a != 0).sum()
@@ -1110,8 +1096,6 @@ def sparse_solve(data, indices, indptr, b):
   )
   @jtu.run_on_devices("cpu", "cuda")
   def test_sparse_qr_linear_solver_grads(self, size, dtype):
-    if jtu.test_device_matches(["cuda"]) and not sptu.GPU_LOWERING_ENABLED:
-      raise unittest.SkipTest('test requires cusparse/cusolver')
     rng = sptu.rand_sparse(self.rng())
     a = rng((size, size), dtype)
     nse = (a != 0).sum()