jax-ml/jax · error · NotImplementedError

JAX Arrays do not implement the arr.flat property: consider

Error message

JAX Arrays do not implement the arr.flat property: consider arr.flatten() instead.

What it means

JAX arrays deliberately do not implement NumPy's `arr.flat` iterator (a lazy 1-d iterator), because lazy iteration conflicts with JAX's tracing/compilation model. A dedicated method raises NotImplementedError with a pointer to the supported alternative. Use flatten() or ravel() for a materialized 1-d view.

Source

Thrown at jax/_src/numpy/array_methods.py:649

    out = _view(self, lax_numpy.finfo(dtype).dtype).astype(dtype)
    return out[..., 0::2] + 1j * out[..., 1::2]

  # lax.bitcast_convert_type adds or subtracts dimensions depending on the
  # relative bitwidths of the dtypes; we account for that with reshapes.
  if nbits_in < nbits_out:
    factor = nbits_out // nbits_in
    out = self.reshape(*self.shape[:-1], self.shape[-1] // factor, factor)
    return lax.bitcast_convert_type(out, dtype)
  elif nbits_in > nbits_out:
    out = lax.bitcast_convert_type(self, dtype)
    return out.reshape(*out.shape[:-2], out.shape[-2] * out.shape[-1])
  else:
    return lax.bitcast_convert_type(self, dtype)


def _notimplemented_flat(self):
  """Not implemented: Use :meth:`~jax.Array.flatten` instead."""
  raise NotImplementedError("JAX Arrays do not implement the arr.flat property: "
                            "consider arr.flatten() instead.")

# TODO(jakevdp): make _accepted_binop_types match the ArrayLike union. Currently
# ArrayLike includes np.number, while here we are more permissive and include
# np.generic: this is required because as of v0.5.X, ml_dtypes types are subclasses
# of np.generic rather than of np.number. Making these match will allow removal of
# cast() calls in the operator definitions below.
_accepted_binop_types = (
    int,
    float,
    complex,
    np.generic,
    np.ndarray,
    Array,
    literals.TypedNdArray,
)

def _operator_eq(self, other):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace arr.flat with arr.flatten() or arr.ravel()
2. For flat indexing use `arr.flatten()[i]` or `arr.take(i)` / `arr unravel_index` logic
3. For iteration over elements, convert to NumPy: `for x in np.asarray(arr).flat` (outside jit)

Example fix

# before
arr.flat[3]

# after
arr.flatten()[3]

flat = arr.flatten()  # instead of arr.flat
first = arr.ravel()[0]

When it happens

Trigger: Accessing `arr.flat` on any jax.Array, including indexing it (`arr.flat[3]`) or iterating (`for x in arr.flat`).

Common situations: Ported NumPy code using arr.flat for iteration or flat indexing; using np.flatnonzero-style idioms on JAX arrays.

Related errors

• Only implemented for order='K'
• np.reshape order=A is not implemented.
• `type` argument of array.view() is not supported.
• array ref with memory space only works inside of a `jit`.
• pinned array ref only works inside of a `jit`.

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