{"record":{"id":"f1d673f7590e8f39","repo":"jax-ml/jax","slug":"jax-arrays-do-not-implement-the-arr-flat-property","errorCode":null,"errorMessage":"JAX Arrays do not implement the arr.flat property: consider arr.flatten() instead.","messagePattern":"JAX Arrays do not implement the arr\\.flat property: consider arr\\.flatten\\(\\) instead\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/array_methods.py","lineNumber":649,"sourceCode":"    out = _view(self, lax_numpy.finfo(dtype).dtype).astype(dtype)\n    return out[..., 0::2] + 1j * out[..., 1::2]\n\n  # lax.bitcast_convert_type adds or subtracts dimensions depending on the\n  # relative bitwidths of the dtypes; we account for that with reshapes.\n  if nbits_in < nbits_out:\n    factor = nbits_out // nbits_in\n    out = self.reshape(*self.shape[:-1], self.shape[-1] // factor, factor)\n    return lax.bitcast_convert_type(out, dtype)\n  elif nbits_in > nbits_out:\n    out = lax.bitcast_convert_type(self, dtype)\n    return out.reshape(*out.shape[:-2], out.shape[-2] * out.shape[-1])\n  else:\n    return lax.bitcast_convert_type(self, dtype)\n\n\ndef _notimplemented_flat(self):\n  \"\"\"Not implemented: Use :meth:`~jax.Array.flatten` instead.\"\"\"\n  raise NotImplementedError(\"JAX Arrays do not implement the arr.flat property: \"\n                            \"consider arr.flatten() instead.\")\n\n# TODO(jakevdp): make _accepted_binop_types match the ArrayLike union. Currently\n# ArrayLike includes np.number, while here we are more permissive and include\n# np.generic: this is required because as of v0.5.X, ml_dtypes types are subclasses\n# of np.generic rather than of np.number. Making these match will allow removal of\n# cast() calls in the operator definitions below.\n_accepted_binop_types = (\n    int,\n    float,\n    complex,\n    np.generic,\n    np.ndarray,\n    Array,\n    literals.TypedNdArray,\n)\n\ndef _operator_eq(self, other):","sourceCodeStart":631,"sourceCodeEnd":667,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/array_methods.py#L631-L667","documentation":"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.","triggerScenarios":"Accessing `arr.flat` on any jax.Array, including indexing it (`arr.flat[3]`) or iterating (`for x in arr.flat`).","commonSituations":"Ported NumPy code using arr.flat for iteration or flat indexing; using np.flatnonzero-style idioms on JAX arrays.","solutions":["Replace arr.flat with arr.flatten() or arr.ravel()","For flat indexing use `arr.flatten()[i]` or `arr.take(i)` / `arr unravel_index` logic","For iteration over elements, convert to NumPy: `for x in np.asarray(arr).flat` (outside jit)"],"exampleFix":"# before\narr.flat[3]\n\n# after\narr.flatten()[3]","handlingStrategy":"type-guard","validationCode":"flat = arr.flatten()  # instead of arr.flat\nfirst = arr.ravel()[0]","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Search-and-replace .flat with .flatten()/[...] when porting","Convert to np.asarray(arr).flat only outside jit"],"tags":["jax","flat","not-implemented","numpy-compat"],"backgroundTag":"unsupported-attribute-access","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}