unskip tests in tests/lax_numpy_test and tests/pallas/gpy_ops_test.py and fix decorator order in tests/linalg_test

Author: gulsumgudukbay

tests/lax_numpy_test.py

@@ -1562,7 +1562,6 @@ def testTrimZerosNotOneDArray(self):
   )
   @jax.default_matmul_precision("float32")
   def testPoly(self, a_shape, dtype, rank):
-    self.skipTest("Skip Poly tests on ROCm")
     if dtype in (np.float16, jnp.bfloat16, np.int16):
       self.skipTest(f"{dtype} gets promoted to {np.float16}, which is not supported.")
     elif rank == 2 and not jtu.test_device_matches(["cpu", "gpu"]):

tests/linalg_test.py

@@ -361,7 +361,7 @@ def testEigvals(self, shape, dtype):
 
   @jtu.run_on_devices("cpu", "gpu")
   def testEigvalsInf(self):
-    self.skipTest("Skip test on ROCm")
+    #self.skipTest("Skip test on ROCm")
     # https://github.com/jax-ml/jax/issues/2661
     if jtu.test_device_matches(["gpu"]) and jtu.jaxlib_version() <= (0, 4, 35):
       self.skipTest("eig on GPU requires jaxlib version > 0.4.35")
@@ -2027,12 +2027,12 @@ def testSqrtmEdgeCase(self, diag, expected, dtype):
 
     self.assertAllClose(root, expected, check_dtypes=False)
 
-  @jtu.ignore_warning(category=FutureWarning, message="Don't treat future SciPy warning as error")
   @jtu.sample_product(
     cshape=[(), (4,), (8,), (4, 7), (2, 1, 5)],
     cdtype=float_types + complex_types,
     rshape=[(), (3,), (7,), (4, 4), (2, 4, 0)],
     rdtype=float_types + complex_types + int_types)
+  @jtu.ignore_warning(category=FutureWarning, message="Don't treat future SciPy warning as error")
   def testToeplitzConstruction(self, rshape, rdtype, cshape, cdtype):
     if ((rdtype in [np.float64, np.complex128]
          or cdtype in [np.float64, np.complex128])

tests/pallas/gpu_ops_test.py

@@ -239,7 +239,6 @@ def test_fused_attention_bwd(
       causal,
       use_segment_ids,
   ):
-    self.skipTest("Skip tests on ROCm")
     k1, k2, k3 = random.split(random.key(0), 3)
     q = random.normal(
         k1, (batch_size, seq_len, num_heads, head_dim), dtype=jnp.float16