[Pallas] Make rng tests inherit from JaxTestCase

PiperOrigin-RevId: 839410599

Author: justinjfu

tests/pallas/tpu_pallas_random_test.py

@@ -243,7 +243,7 @@ def f(rng_key):
     self.assertGreaterEqual(jnp.max(y), jnp.min(y))
 
 
-class BlockInvarianceTest(parameterized.TestCase):
+class BlockInvarianceTest(jtu.JaxTestCase):
 
   def setUp(self):
     if not jtu.test_device_matches(["tpu"]):
@@ -290,7 +290,7 @@ def body(key_ref, o_ref):
     np.testing.assert_array_equal(result_16x128, result_32x256)
 
 
-class ThreefryTest(parameterized.TestCase):
+class ThreefryTest(jtu.JaxTestCase):
 
   def setUp(self):
     if not jtu.test_device_matches(["tpu"]):
@@ -373,7 +373,7 @@ def test_threefry_kernel_matches_jax_threefry_sharded(self, shape):
     np.testing.assert_array_equal(jax_gen, pl_gen)
 
 
-class PhiloxTest(parameterized.TestCase):
+class PhiloxTest(jtu.JaxTestCase):
 
   def setUp(self):
     if not jtu.test_device_matches(["tpu"]):