Remove code that skips array PRNG tests on CUDA.

#13037 added this skip, but I have no idea why. The test seems to pass on GPU. PiperOrigin-RevId: 568216977
google · Sep 25, 2023 · d3f5e7f · d3f5e7f
1 parent 3a9289e
commit d3f5e7f
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/tests/array_test.py b/tests/array_test.py
@@ -1125,7 +1125,6 @@ class RngShardingTest(jtu.JaxTestCase):
   # tests that the PRNGs are automatically sharded as expected
 
   @parameterized.named_parameters(("3", 3), ("4", 4), ("5", 5))
-  @jtu.skip_on_devices("gpu")
   def test_random_bits_is_pure_map_1d(self, num_devices):
     @jax.jit
     def f(x):
@@ -1159,7 +1158,6 @@ def f(x):
        "mesh_shape": mesh_shape, "pspec": pspec}
       for mesh_shape in [(3, 2), (4, 2), (2, 3)]
       for pspec in [P('x', None), P(None, 'y'), P('x', 'y')])
-  @jtu.skip_on_devices("gpu")
   def test_random_bits_is_pure_map_2d(self, mesh_shape, pspec):
     @jax.jit
     def f(x):