jax-ml/jax · error · TypeError

iteration over a 0-d array

Error message

iteration over a 0-d array

What it means

jax.Array.__iter__ raises TypeError when iterating over a 0-d (scalar) array, mirroring NumPy's behavior since scalars are not iterable. JAX implements the Python iteration protocol only for arrays with ndim >= 1. The error is raised before any sharding/addressability checks.

Source

Thrown at jax/_src/array.py:346

  def __format__(self, format_spec):
    if isinstance(self.sharding, NamedSharding) and self.sharding.spec.unreduced:
      return repr(self)
    elif (self.is_fully_addressable or self.is_fully_replicated and
          self.sharding.has_addressable_devices):
      # Simulates behavior of https://github.com/numpy/numpy/pull/9883
      return format(self._value if self.ndim else self._value[()], format_spec)
    else:
      return repr(self)

  def __getitem__(self, idx, /):
    from jax._src.numpy import indexing  # pyrefly: ignore[missing-import]
    self._check_if_deleted()

    return indexing.rewriting_take(self, idx)

  def __iter__(self):
    if self.ndim == 0:
      raise TypeError("iteration over a 0-d array")  # same as numpy error
    else:
      assert self.is_fully_replicated or self.is_fully_addressable
      if self.sharding.num_devices == 1 or self.is_fully_replicated:
        return (sl for chunk in self._chunk_iter(100) for sl in chunk._unstack())  # pyrefly: ignore[missing-attribute]
      else:
        # TODO(yashkatariya): Don't bounce to host and use `_chunk_iter` path
        # here after uneven partitioning support is added.
        return (api.device_put(self._value[i]) for i in range(self.shape[0]))

  @property
  def is_fully_replicated(self) -> bool:
    return self.sharding.is_fully_replicated

  def __repr__(self):
    prefix = 'Array('
    if self.aval is not None and self.aval.weak_type:
      dtype_str = f'dtype={self.dtype.name}, weak_type=True'
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use x.item() to extract the Python scalar: float(x.item())
2. Reshape to 1-d before iterating: x.reshape(1) or x[None]
3. Check rank first: if x.ndim == 0: handle scalar path
4. Use jnp.atleast_1d(x) before generic iteration code

Example fix

// before
loss = jnp.mean(preds - targets)
for v in loss:  # TypeError: iteration over a 0-d array
  ...
// after
loss = jnp.mean(preds - targets)
val = loss.item()
# or iterate a 1-d view:
for v in jnp.atleast_1d(loss):
  ...

def is_scalar_array(x):
    return hasattr(x, 'ndim') and x.ndim == 0

if is_scalar_array(loss):
    total = loss.item()
else:
    total = float(jnp.sum(loss))

from jax import Array
import numpy as np

def is_zero_dim(x) -> bool:
    """True for jax arrays, ndarrays, or scalars with ndim == 0."""
    return isinstance(x, (Array, np.ndarray)) and getattr(x, 'ndim', -1) == 0

try:
    for v in maybe_scalar:
        process(v)
except TypeError as e:
    if 'iteration over a 0-d array' in str(e):
        process(maybe_scalar.item())
    else:
        raise

When it happens

Trigger: Calling list(x), for v in x, tuple(x), or *x unpacking on an Array with x.ndim == 0, e.g. the result of jnp.squeeze, .sum(), .mean(), or indexing a 1-d array with a single int.

Common situations: Computing a scalar loss/accuracy and accidentally iterating it; aggressive jnp.squeeze removing all dims; functions that accept 'array or scalar' and do `for item in arg`; switching code from Python floats (iterable via string repr confusion) or assuming jnp scalars behave like 1-element lists.

Related errors

• Formatting arguments to checkify.check need to be PyTrees of
• check_error takes an Error as argument, got type {type(error
• bool() not supported for instances of type '{0}' (did you me
• Default value must be of type bool, got {default} of type {g
• Default value must be of type str, got {default} of type {ge

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