skip sparse_nm tests on ROCm

Author: Ruturaj4

tests/sparse_nm_test.py

@@ -47,6 +47,8 @@ def setUp(self):
   )
   @jtu.run_on_devices("gpu")
   def test_shapes(self, tile_m, tile_n, tile_k, batch, sparse_idx):
+    if jtu.is_device_rocm():
+        self.skipTest("There is no spmm lowering support for ROCm platform")
     # Build keyword arguments
     kwargs = {
         "dimension_numbers": (((1,), (1,)), (tuple(), tuple())),
@@ -93,6 +95,8 @@ def test_shapes(self, tile_m, tile_n, tile_k, batch, sparse_idx):
   )
   @jtu.run_on_devices("gpu")
   def test_types(self, lhs_type, rhs_type, output_type):
+    if jtu.is_device_rocm():
+        self.skipTest("There is no spmm lowering support for ROCm platform")
     tile_m, tile_n, tile_k = 64, 32, 128
 
     # Build input data