{"record":{"id":"5ee7ce1525f48ef8","repo":"jax-ml/jax","slug":"iteration-over-a-0-d-key-array","errorCode":null,"errorMessage":"iteration over a 0-d key array","messagePattern":"iteration over a 0-d key array","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/prng.py","lineNumber":298,"sourceCode":"  def sharding(self):\n    return logical_sharding(self.shape, self.dtype, self._base_array.sharding)\n\n  @property\n  def committed(self):\n    return self._base_array.committed\n\n  def _is_scalar(self):\n    base_ndim = len(self._impl.key_shape)\n    return self._base_array.ndim == base_ndim\n\n  def __len__(self):\n    if self._is_scalar():\n      raise TypeError('len() of unsized object')\n    return len(self._base_array)\n\n  def __iter__(self) -> Iterator[PRNGKeyArray]:\n    if self._is_scalar():\n      raise TypeError('iteration over a 0-d key array')\n    # TODO(frostig): we may want to avoid iteration by slicing because\n    # a very common use of iteration is `k1, k2 = split(key)`, and\n    # slicing/indexing may be trickier to track for linearity checking\n    # purposes. Maybe we can:\n    # * introduce an unpack primitive+traceable (also allow direct use)\n    # * unpack upfront into shape[0] many keyarray slices\n    # * return iter over these unpacked slices\n    # Whatever we do, we'll want to do it by overriding\n    # ShapedArray._iter when the element type is KeyTy...\n    return (PRNGKeyArray(self._impl, k) for k in iter(self._base_array))\n\n  def __bool__(self):\n    raise TypeError(\"key array cannot be converted to boolean.\")\n\n  def __repr__(self):\n    return (f'Array({self.shape}, dtype={self.dtype.name}) overlaying:\\n'\n            f'{self._base_array}')\n","sourceCodeStart":280,"sourceCodeEnd":316,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/prng.py#L280-L316","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Split first: k1, k2 = jax.random.split(key, 2)","Check key.ndim/key.shape before iterating","Return an explicit list/tuple of keys from APIs instead of a single key when callers unpack"],"exampleFix":"// before\nk1, k2 = key  # TypeError: iteration over a 0-d key array\n\n// after\nk1, k2 = jax.random.split(key, 2)","handlingStrategy":"type-guard","validationCode":"if key.ndim == 0:\n    raise ValueError('single key; split before iterating')\nkeys = list(key)","typeGuard":"def is_iterable_key(key) -> bool:\n    return key.ndim > 0 and not (key.ndim == len(key._impl.key_shape)) if hasattr(key, '_impl') else key.ndim > 0","tryCatchPattern":null,"preventionTips":["Use jax.random.split(key, n) to get iterable key batches","Unpack only split() results, never bare keys"],"tags":["jax","prng","iteration","unpacking"],"backgroundTag":"jax-prng-key-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}