jax-ml/jax · error · NotImplementedError

jnp.fromiter() is not implemented because it may be non-pure

Error message

jnp.fromiter() is not implemented because it may be non-pure and thus unsafe for use with JIT and other JAX transformations. Consider using jnp.asarray(np.fromiter(...)) instead, although care should be taken if np.fromiter is used within a jax transformations because of its potential side-effect of consuming the iterable object; for more information see https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions

What it means

jnp.fromiter is unimplemented for the same purity reason as jnp.fromfile: consuming an iterator is a side effect incompatible with JIT and JAX transformations. JAX requires array construction to be deterministic and side-effect free, so the message points to np.fromiter plus jnp.asarray.

Source

Thrown at jax/_src/numpy/lax_numpy.py:5551

    "jnp.fromfile() is not implemented because it may be non-pure and thus unsafe for use "
    "with JIT and other JAX transformations. Consider using jnp.asarray(np.fromfile(...)) "
    "instead, although care should be taken if np.fromfile is used within a jax transformations "
    "because of its potential side-effect of consuming the file object; for more information see "
    "https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions")


@export
def fromiter(*args, **kwargs):
  """Unimplemented JAX wrapper for jnp.fromiter.

  This function is left deliberately unimplemented because it may be non-pure and thus
  unsafe for use with JIT and other JAX transformations. Consider using
  ``jnp.asarray(np.fromiter(...))`` instead, although care should be taken if ``np.fromiter``
  is used within jax transformations because of its potential side-effect of consuming the
  iterable object; for more information see `Common Gotchas: Pure Functions
  <https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions>`_.
  """
  raise NotImplementedError(
    "jnp.fromiter() is not implemented because it may be non-pure and thus unsafe for use "
    "with JIT and other JAX transformations. Consider using jnp.asarray(np.fromiter(...)) "
    "instead, although care should be taken if np.fromiter is used within a jax transformations "
    "because of its potential side-effect of consuming the iterable object; for more information see "
    "https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions")


@export
def from_dlpack(x: Any, /, *, device: xc.Device | Sharding | None = None,
                copy: bool | None = None) -> Array:
  """Construct a JAX array via DLPack.

  JAX implementation of :func:`numpy.from_dlpack`.

  Args:
    x: An object that implements the DLPack_ protocol via the ``__dlpack__``
      and ``__dlpack_device__`` methods, or a legacy DLPack tensor on either
      CPU or GPU.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use jnp.asarray(np.fromiter(it, dtype, count))
Or materialize the iterable first: jnp.asarray(list(it), dtype=dtype)
Do the conversion outside of any jitted computation

Example fix

// before
x = jnp.fromiter(gen, dtype=np.float32)
// after
x = jnp.asarray(np.fromiter(gen, dtype=np.float32))

Defensive patterns

Strategy: fallback

Prevention

Grep for jnp.fromiter when porting; replace with jnp.asarray(np.fromiter(...))
Materialize iterables to lists/arrays outside jit boundaries

When it happens

Trigger: Any call to jnp.fromiter(iterable, dtype) — unconditionally raises NotImplementedError. Typical when porting code that streams values (generators, database cursors) into an array.

Common situations: Generator-based data ingestion code ported from NumPy; building arrays lazily inside functions later wrapped in jax.jit.

Related errors

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