From 1fec09c166e97fb72e5664161f93cbf5e7508dc4 Mon Sep 17 00:00:00 2001 From: Maddipatla Chatan Date: Fri, 18 Sep 2026 06:24:19 +0530 Subject: [PATCH] Handle JAX_MODE for valid PRNG key generation Added handling for JAX_MODE to return a valid JAX PRNG key instead of a plain integer seed. --- tensorflow_probability/python/util/seed_stream.py | 10 ++++++++++ 1 file changed, 10 insertions(+) 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