Fix test_convert_seed.

PiperOrigin-RevId: 441138590
google-deepmind · Apr 12, 2022 · ad5e452 · ad5e452
1 parent 95056e2
commit ad5e452
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/distrax/_src/distributions/distribution_test.py b/distrax/_src/distributions/distribution_test.py
@@ -129,7 +129,7 @@ def sampler(dist, seed):
   )
   def test_convert_seed(self, dtype):
     rng, _ = distribution.convert_seed_and_sample_shape(dtype(0), 2)
-    assert isinstance(rng, distribution.PRNGKey)
+    jax.random.split(rng)  # Should not raise an error.
 
   @parameterized.named_parameters(
       ('int', 2, (2,)),