Skip test_batch_axis_sharding_jvp because of hipSPARSE issue

Author: AratiGanesh

tests/linalg_sharding_test.py

@@ -158,6 +158,8 @@ def test_non_batch_axis_sharding(self, fun_and_shapes, dtype):
   )
   @jtu.run_on_devices("gpu", "cpu")
   def test_batch_axis_sharding_jvp(self, fun_and_shapes, dtype):
+    if fun_and_shapes[0] is lax.linalg.tridiagonal_solve and jtu.is_device_rocm():
+      self.skipTest("test_batch_axis_sharding_jvp is not supported on ROCm")
     fun, shapes = self.get_fun_and_shapes(fun_and_shapes, grad=True)
     primals = self.get_args(shapes, dtype, batch_size=8)
     tangents = tuple(map(jnp.ones_like, primals))