[Pallas] Make rng tests inherit from JaxTestCase

PiperOrigin-RevId: 839963758

Author: justinjfu

jax/experimental/pallas/ops/tpu/random/philox.py

@@ -117,7 +117,7 @@ def kernel(offset_ref, key_ref, out_ref):
     offset = prng_utils.compute_scalar_offset(
         counts_idx, unpadded_shape, block_shape)
     counts_lo = prng_utils.blocked_iota(block_size, unpadded_shape)
-    counts_lo = counts_lo + offset + offset_ref[0]
+    counts_lo = counts_lo + offset.astype(jnp.uint32) + offset_ref[0]
     counts_lo = counts_lo.astype(jnp.uint32)
     # TODO(justinfu): Support hi bits on count.
     _zeros = jnp.zeros_like(counts_lo)

jax/experimental/pallas/ops/tpu/random/threefry.py

@@ -63,7 +63,7 @@ def kernel(key_ref, out_ref):
     offset = prng_utils.compute_scalar_offset(
         counts_idx, unpadded_shape, block_shape)
     counts_lo = prng_utils.blocked_iota(block_size, unpadded_shape)
-    counts_lo = counts_lo + offset
+    counts_lo = counts_lo + offset.astype(jnp.uint32)
     counts_lo = counts_lo.astype(jnp.uint32)
     # TODO(justinfu): Support hi bits on count.
     counts_hi = jnp.zeros_like(counts_lo)

tests/pallas/tpu_pallas_random_test.py

@@ -197,7 +197,7 @@ def body(key_ref, o_ref):
           key, shape=o_ref[0, ...].shape, minval=0.0, maxval=1.0
       )
 
-      key = jax_random.fold_in(key, 2)
+      key = jax_random.fold_in(key, jnp.uint32(2))
       o_ref[1, ...] = jax_random.uniform(
           key, shape=o_ref[1, ...].shape, minval=0.0, maxval=1.0
       )
@@ -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"]):