diff --git a/tensorflow_probability/python/util/seed_stream.py b/tensorflow_probability/python/util/seed_stream.py index c90725e2ce..0f0ca1e106 100644 --- a/tensorflow_probability/python/util/seed_stream.py +++ b/tensorflow_probability/python/util/seed_stream.py @@ -206,6 +206,16 @@ def __call__(self): seed: A fresh integer usable as a seed in downstream operations, or `None`. """ + # Under JAX_MODE, `_seed` has already been converted (in `__init__`) into + # a real `jax.random.PRNGKey`, not a plain int/None -- so this must + # dispatch to `_call_jax`, which correctly derives a new key via + # `jaxrand.fold_in`, rather than falling through to the SHA512-of-a- + # string-repr path below, which was previously reachable even under + # JAX_MODE and returns a plain Python int that is not a valid JAX PRNG + # key (confirmed: passing that int straight to e.g. `jax.random.normal` + # raises `TypeError: JAX encountered invalid PRNG key data`). + if JAX_MODE: + return self._call_jax() self._counter += 1 if self._seed is None: return None