jax-ml/jax · error · ValueError

Expected base_key to be a typed PRNG key; got {self._base_ke

Error message

Expected base_key to be a typed PRNG key; got {self._base_key}

What it means

jax.experimental.random.stateful_rng() optionally accepts a base_key (a typed JAX PRNG key, dtype key<fry>). In __post_init__ it validates that base_key is a JAX Array whose dtype is a prng_key subtype; passing a legacy uint32 PRNGKey, a seed integer, or any other array raises this ValueError.

Source

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

    _counter: a scalar integer wrapped in a :class:`jax.Ref`.

  Examples:

  >>> from jax.experimental import random
  >>> rng = random.stateful_rng(42)
  >>> rng
  StatefulPRNG(_base_key=Array((), dtype=key<fry>) overlaying:
  [ 0 42], _counter=Ref(0, dtype=int32, weak_type=True))
  """
  _base_key: Array
  _counter: core.Ref

  def __post_init__(self):
    if self._base_key is api_util.SENTINEL:
      return
    if not (isinstance(self._base_key, Array)
            and dtypes.issubdtype(self._base_key.dtype, dtypes.prng_key)):
      raise ValueError(f"Expected base_key to be a typed PRNG key; got {self._base_key}")

    # TODO(jakevdp): how to validate a traced mutable array?
    if not (isinstance(self._counter, core.Ref) or
            (isinstance(self._counter, core.Tracer)
             and isinstance(self._counter.aval, state_types.AbstractRef))):
      raise ValueError(f"Expected counter to be a scalar integer ref; got {self._counter}")

  def key(self, shape: int | Sequence[int] = ()) -> Array:
    """Generate a new JAX PRNGKey, updating the internal state.

    Args:
      shape: an optional shape if returning multiple keys.

    Returns:
      A new, independent PRNG key with the same impl/dtype as
      ``self._base_key``.

    Examples:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert to a typed key first: stateful_rng(base_key=jax.random.key(0))
2. If you only have a seed, pass stateful_rng(seed=0) instead
3. If you have raw data, wrap it: jax.random.random_wrap(data, impl='threefry2x32')

Example fix

// before
rng = stateful_rng(base_key=jax.random.PRNGKey(0))

// after
rng = stateful_rng(base_key=jax.random.key(0))

import jax, jax.numpy as jnp
if not (jax.dtypes.issubdtype(base_key.dtype, jax.dtypes.prng_key)):
    base_key = jax.random.key(0)  # or convert appropriately

import jax
def is_typed_key(x) -> bool:
    return isinstance(x, jax.Array) and jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key)

When it happens

Trigger: Calling stateful_rng(base_key=jax.random.PRNGKey(0)) (legacy uint32 key); passing a Python int seed as base_key instead of the seed= parameter; passing a raw uint32 array produced by key_data.

Common situations: Migrating old code that held jax.random.PRNGKey objects to the stateful API; parameter mix-ups between seed= and base_key=; loading unwrapped key bits from a checkpoint.

Related errors

• Expected counter to be a scalar integer ref; got {self._coun
• cannot operate on split stateful generator
• When used within transformed code, jax.experimental.random.s
• PRNG keys must be loaded from SMEM. Did you set the memory s
• Seed key_data must be 1D.

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