Repository: google/jax

Issue:
Support for `numpy.fromiter`
As far as I can tell there is not an existing issue or implementation for `jax.numpy.fromiter`, which currently raises:

```
[NotImplementedError: Numpy function <built-in function fromiter> not yet implemented]()
```


Discussion:
Thanks for the issue! We haven't implemented this because there's not really any way to do better than `jnp.asarray(np.fromiter(x))`, though we could certainly add that function as a convenience. What do you think?
That's exactly what I ended up doing – I think it would be neat to add it just to get more function parity with `numpy`, even if it doesn't really do anything clever.
@jakevdp Jax could actually do something 'clever' for `fromiter(x)`, as `jnp.asarray(np.fromiter(x))` materelize everything eagerly. That said, the staging behaviour can be achieved with something like `jnp.array(list(x))`.
Initial implementations of `fromiter` and related functions in #10049. Thanks!
Wow! Very impressive turn-around time on these issues! Thank you so much @jakevdp!
Some of the functions in this family—namely `fromiter` and `fromfile`—are inherently side-effecting and thus incompatible with JAX staging (`jit` etc.). I think we should avoid introducing those into `jax.numpy`. Parity with standard numpy is only one goal; others include behavior invariance under `jit` and avoiding impure functions in our core libraries.

Other functions seem like a plus. For instance, in the context of staging, `fromstring` could be a convenient way to introduce a constant.
Thanks Roy - good points! I'll make the modifications to #10049
So the resolution of the original question is that we don't want to support `jnp.fromiter` in JAX because it's a potentially side-effecting operation that will possibly cause surprising behavior with JIT. But we were able to implement a few other `jnp.from*` operations.

Thanks for the request!
Might it be worthwhile to add a description to the `jnp.fromiter` function such that it explains the reasoning @froystig mentioned in the error instead of `Numpy function <built-in function fromiter> not yet implemented]()` in addition to a suggested solution (`jnp.asarray(np.fromiter(x))` along with impure + jit caveat)?

One thing I really like about JAX is when the errors have references to solutions and rationale, which the library does excellently in many other places!
That's a great idea!

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1973,6 +1973,42 @@ def frombuffer(buffer, dtype=float, count=-1, offset=0):
   return asarray(np.frombuffer(buffer=buffer, dtype=dtype, count=count, offset=offset))
 
 
+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://jax.readthedocs.io/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://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions")
+
+
+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://jax.readthedocs.io/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://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions")
+
+
 @_wraps(np.fromfunction)
 def fromfunction(function, shape, *, dtype=float, **kwargs):
   shape = core.canonicalize_shape(shape, context="shape argument of jnp.fromfunction()")
diff --git a/jax/numpy/__init__.py b/jax/numpy/__init__.py
--- a/jax/numpy/__init__.py
+++ b/jax/numpy/__init__.py
@@ -123,7 +123,9 @@
     fmax as fmax,
     fmin as fmin,
     frombuffer as frombuffer,
+    fromfile as fromfile,
     fromfunction as fromfunction,
+    fromiter as fromiter,
     fromstring as fromstring,
     full as full,
     full_like as full_like,


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -6245,7 +6245,7 @@ def test_lax_numpy_docstrings(self):
     # Test that docstring wrapping & transformation didn't fail.
 
     # Functions that have their own docstrings & don't wrap numpy.
-    known_exceptions = {'broadcast_arrays', 'vectorize'}
+    known_exceptions = {'broadcast_arrays', 'fromfile', 'fromiter', 'vectorize'}
 
     for name in dir(jnp):
       if name in known_exceptions or name.startswith('_'):