jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

jnp.fromfile is intentionally not implemented in JAX because reading from a file object is a side effect (it consumes the file), violating the purity requirements of JIT and other transformations. The error message directs users to load data with numpy first and convert via jnp.asarray.

Source

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

    Array([0, 1, 2, 3, 4], dtype=int32)

  .. _Python buffer interface: https://docs.python.org/3/c-api/buffer.html
  """
  return asarray(np.frombuffer(buffer=buffer, dtype=dtype, count=count, offset=offset))


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

  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.fromfile(...))`` instead, although care should be taken if ``np.fromfile``
  is used within jax transformations because of its potential side-effect of consuming the
  file 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.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>`_.
  """

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Replace with jnp.asarray(np.fromfile(path, dtype))
Better: use np.load / np.memmap for .npy files, then jnp.asarray
Ensure loading happens outside jax.jit / pmap and other transformations

Example fix

// before
x = jnp.fromfile('data.bin', dtype=np.float32)
// after
x = jnp.asarray(np.fromfile('data.bin', dtype=np.float32))

Defensive patterns

Strategy: fallback

Prevention

Grep for jnp.fromfile when upgrading/porting; replace with jnp.asarray(np.fromfile(...))
Keep all file I/O outside jitted functions
Use np.load/np.memmap for structured files

When it happens

Trigger: Any call to jnp.fromfile(file, dtype=...) — it unconditionally raises NotImplementedError. Commonly hit when porting NumPy code that loads binary data (e.g. MNIST images, .npy raw dumps) directly.

Common situations: Ported np.fromfile pipelines for binary datasets; attempting to load data inside a jitted function or during sharded initialization where the file read would be replicated and impure.

Related errors

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