jax-ml/jax · error · ValueError

cannot operate on split stateful generator

Error message

cannot operate on split stateful generator

What it means

jax.experimental.random.stateful_rng() supports splitting the generator so children draw independent streams, but a split generator's _base_key has a non-empty shape. Operations that need to advance the scalar counter (key(), random(), uniform(), normal(), integers(), split(), spawn()) refuse to run on such a split generator via this ValueError.

Source

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

      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()
      Array((), dtype=key<fry>) overlaying:
      [ 928981903 3453687069]
    """
    if self._base_key.shape:
      # TODO(jakevdp): better error message.
      raise ValueError("cannot operate on split stateful generator")

    key = random.fold_in(self._base_key, ref_primitives.ref_get(self._counter))
    ref_primitives.ref_addupdate(self._counter, ..., 1)
    shape_tuple = _canonicalize_size(shape)
    return random.split(key, shape_tuple) if shape_tuple else key

  def random(
      self,
      size: int | Sequence[int] | None = None,
      dtype: DTypeLike = float,
  ):
    """Return random floats in the half-open interval [0.0, 1.0)."""
    # TODO(jakevdp): write docstring
    return random.uniform(self.key(), shape=_canonicalize_size(size), dtype=dtype)


  def uniform(
      self,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Only call draw methods on the top-level generator; use spawned children to spawn further or pass as base_key to new stateful_rng instances
2. Use rng.spawn() at the point where independent substreams are needed instead of pre-splitting
3. Restructure so each worker receives its own stateful_rng(seed=unique_seed)

Example fix

// before
children = rng.split(2)
k = children[0].key()  # ValueError: cannot operate on split stateful generator

// after
child = rng.spawn(1)[0]
rng2 = stateful_rng(base_key=child.key())
k = rng2.key()

import jax
if jax.random.key_data(rng._base_key).shape != (2,):
    raise ValueError('generator is split; spawn children instead of drawing')

def is_unsplit_generator(rng) -> bool:
    return rng._base_key.shape == ()

When it happens

Trigger: Calling rng.key() or any sampling method on an object obtained from rng.spawn() or rng.split(); reusing a child generator for direct draws instead of spawning further children.

Common situations: Adopting the stateful API and assuming split children behave like independent top-level generators; passing spawned generators into code that calls rng.random() directly.

Related errors

• Expected base_key to be a typed PRNG key; got {self._base_ke
• Expected counter to be a scalar integer ref; got {self._coun
• When used within transformed code, jax.experimental.random.s
• Sizes passed to split must be nonnegative, got {list(sizes)}
• Sum of sizes {np.sum(sizes)} must be equal to dimension {axi

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