jax-ml/jax · error · ValueError

Expected counter to be a scalar integer ref; got {self._coun

Error message

Expected counter to be a scalar integer ref; got {self._counter}

What it means

The stateful RNG in jax.experimental.random tracks draw counts in a counter that must be a mutable scalar integer ref (core.Ref, or a tracer with an AbstractRef aval), enabling safe state updates under jit. __post_init__ raises this ValueError when counter is a plain array, Python int, or anything else non-ref.

Source

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

  >>> 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:
      >>> from jax.experimental import random
      >>> rng = random.stateful_rng(0)
      >>> rng.key()
      Array((), dtype=key<fry>) overlaying:
      [1797259609 2579123966]
      >>> rng.key()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Do not construct the state object directly — call stateful_rng(seed=...) which builds a valid counter internally
2. If restoring state, re-derive the counter via the public API rather than serializing it
3. Check jax version compatibility for the experimental stateful module

Example fix

// before
rng = stateful_rng(seed=0)
rng._counter = jnp.int32(0)  # later use raises

// after
rng = stateful_rng(seed=0)  # keep internal counter untouched

# Users cannot easily construct a valid Ref counter; rely on the public API
rng = stateful_rng(seed=0)  # constructs a well-formed counter internally

When it happens

Trigger: Manually constructing the stateful-RNG state with counter=np.int64(0) or a jnp scalar; passing a regular JAX array as counter; refactoring internals where the counter object loses its Ref type.

Common situations: Users poking at the experimental stateful API internals instead of the public constructor; partial checkpointing/restoration of state; version changes in the experimental ref-state API surface.

Related errors

• Expected base_key to be a typed PRNG key; got {self._base_ke
• cannot operate on split stateful generator
• When used within transformed code, jax.experimental.random.s
• State effect not supported in vmap-of-cond.
• Effects not supported in `scan`: {disallowed_effects}

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