jax-ml/jax · error · TypeError

iteration over a 0-d key array

Error message

iteration over a 0-d key array

What it means

PRNGKeyArray.__iter__ refuses to iterate a scalar (0-d) key because there is no batch axis to iterate over, mirroring the behavior of 0-d JAX arrays. Iteration is only meaningful for a batch of keys, e.g. k1, k2 = jax.random.split(key). The check uses _is_scalar(): base array ndim equals the impl key_shape length.

Source

Thrown at jax/_src/random/prng.py:298

  def sharding(self):
    return logical_sharding(self.shape, self.dtype, self._base_array.sharding)

  @property
  def committed(self):
    return self._base_array.committed

  def _is_scalar(self):
    base_ndim = len(self._impl.key_shape)
    return self._base_array.ndim == base_ndim

  def __len__(self):
    if self._is_scalar():
      raise TypeError('len() of unsized object')
    return len(self._base_array)

  def __iter__(self) -> Iterator[PRNGKeyArray]:
    if self._is_scalar():
      raise TypeError('iteration over a 0-d key array')
    # TODO(frostig): we may want to avoid iteration by slicing because
    # a very common use of iteration is `k1, k2 = split(key)`, and
    # slicing/indexing may be trickier to track for linearity checking
    # purposes. Maybe we can:
    # * introduce an unpack primitive+traceable (also allow direct use)
    # * unpack upfront into shape[0] many keyarray slices
    # * return iter over these unpacked slices
    # Whatever we do, we'll want to do it by overriding
    # ShapedArray._iter when the element type is KeyTy...
    return (PRNGKeyArray(self._impl, k) for k in iter(self._base_array))

  def __bool__(self):
    raise TypeError("key array cannot be converted to boolean.")

  def __repr__(self):
    return (f'Array({self.shape}, dtype={self.dtype.name}) overlaying:\n'
            f'{self._base_array}')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split first: k1, k2 = jax.random.split(key, 2)
2. Check key.ndim/key.shape before iterating
3. Return an explicit list/tuple of keys from APIs instead of a single key when callers unpack

Example fix

// before
k1, k2 = key  # TypeError: iteration over a 0-d key array

// after
k1, k2 = jax.random.split(key, 2)

if key.ndim == 0:
    raise ValueError('single key; split before iterating')
keys = list(key)

def is_iterable_key(key) -> bool:
    return key.ndim > 0 and not (key.ndim == len(key._impl.key_shape)) if hasattr(key, '_impl') else key.ndim > 0

When it happens

Trigger: Unpacking a single key: a, b = key; passing a scalar key to code that loops over key batches; using list(key) or tuple unpacking on jax.random.fold_in output.

Common situations: Refactoring where split(key) was replaced by fold_in or a bare key; generic Python code that does 'for k in keys' over user-supplied values; tuple-unpacking in function returns that used to contain split keys.

Related errors

• iteration over a 0-d array
• iteration over a 0-d array
• `linearize_from_jvp` is a pair of rules, not a single rule;
• `vjp_from_jvp` is a pair of rules, not a single rule; unpack
• `vjp_from_lin` is a pair of rules, not a single rule; unpack

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