jax-ml/jax · error · TypeError

When used within transformed code, jax.experimental.random.s

Error message

When used within transformed code, jax.experimental.random.stateful_rng() requires an explicit seed to be set.

What it means

stateful_rng() with no seed normally derives a random seed from OS entropy at top level. Inside a JAX transformation (jit, grad, vmap, pmap, scan, etc.) that is disallowed — trace-time state creation with implicit entropy would break tracing and reproducibility — so it raises this TypeError unless seed is passed explicitly.

Source

Thrown at jax/_src/random/stateful_rng.py:292

    >>> import jax
    >>> jit_uniform = jax.jit(rng.uniform)
    >>> jit_uniform()
    Array(0.6672406, dtype=float32)
    >>> jit_uniform()
    Array(0.3890121, dtype=float32)

    Keys can be generated directly if desired:

    >>> rng.key()
    Array((), dtype=key<fry>) overlaying:
    [2954079971 3276725750]
    >>> rng.key()
    Array((), dtype=key<fry>) overlaying:
    [2765691542  824333390]
  """
  if seed is None:
    if not core.trace_ctx.is_top_level():
      raise TypeError(
        "When used within transformed code, jax.experimental.random.stateful_rng()"
        " requires an explicit seed to be set.")
    entropy = np.random.SeedSequence().entropy
    assert isinstance(entropy, int)
    seed = np.int64(entropy & np.iinfo(np.int64).max)
  assert seed is not None
  return StatefulPRNG(
    _base_key=random.key(seed, impl=impl),
    _counter=ref.new_ref(0)
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create the generator outside the transformation and pass it in as an argument or closure
2. Pass an explicit seed: stateful_rng(seed=some_int) when creation must happen inside
3. Use spawn()/split() on an outer generator to derive child keys inside traced code

Example fix

// before
@jax.jit
def f(x):
    rng = stateful_rng()  # TypeError inside transformed code
    return x + rng.random()

// after
@jax.jit
def f(x, rng):
    return x + rng.random()
rng = stateful_rng(seed=0)
y = f(x, rng)

import jax
def make_rng(seed=None):
    if seed is None and not jax.core.trace_ctx.is_top_level():
        raise ValueError('pass an explicit seed inside transformed code')
    return stateful_rng(seed)

When it happens

Trigger: Calling stateful_rng() inside a @jax.jit function, in a grad/vmap/scan body, or during any non-top-level trace; constructing generators lazily inside model init functions that get jitted.

Common situations: Model constructors run under hk.transform/jit/Flax NNX where the code creates an RNG without threading a seed in; refactoring top-level generator creation into helper functions called from traced code.

Related errors

• The unsafe_buffer_pointer() method was called on {self._erro
• function {dbg.func_src_info} traced for {dbg.traced_for} ret
• Array slice indices must have static start/stop/step to be u
• Expected base_key to be a typed PRNG key; got {self._base_ke
• Expected counter to be a scalar integer ref; got {self._coun

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/35a62a6a471d2441. Report an issue: GitHub.