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

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
@@ -1966,6 +1966,27 @@ def array_equiv(a1, a2):
 empty = zeros
 
 
+# General np.from* style functions mostly delegate to numpy.
+
+@_wraps(np.frombuffer)
+def frombuffer(buffer, dtype=float, count=-1, offset=0):
+  return asarray(np.frombuffer(buffer=buffer, dtype=dtype, count=count, offset=offset))
+
+
+@_wraps(np.fromfunction)
+def fromfunction(function, shape, *, dtype=float, **kwargs):
+  shape = core.canonicalize_shape(shape, context="shape argument of jnp.fromfunction()")
+  for i in range(len(shape)):
+    in_axes = [0 if i == j else None for j in range(len(shape))]
+    function = jax.vmap(function, in_axes=tuple(in_axes[::-1]))
+  return function(*(arange(s, dtype=dtype) for s in shape), **kwargs)
+
+
+@_wraps(np.fromstring)
+def fromstring(string, dtype=float, count=-1, *, sep):
+  return asarray(np.fromstring(string=string, dtype=dtype, count=count, sep=sep))
+
+
 @_wraps(np.eye)
 def eye(N, M=None, k=0, dtype=None):
   lax_internal._check_user_dtype_supported(dtype, "eye")
diff --git a/jax/numpy/__init__.py b/jax/numpy/__init__.py
--- a/jax/numpy/__init__.py
+++ b/jax/numpy/__init__.py
@@ -122,6 +122,9 @@
     floating as floating,
     fmax as fmax,
     fmin as fmin,
+    frombuffer as frombuffer,
+    fromfunction as fromfunction,
+    fromstring as fromstring,
     full as full,
     full_like as full_like,
     gcd as gcd,


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
@@ -5922,6 +5922,27 @@ def testDefaultDtypes(self):
     self.assertEqual(jnp.float_, np.float32 if precision == '32' else np.float64)
     self.assertEqual(jnp.complex_, np.complex64 if precision == '32' else np.complex128)
 
+  def testFromBuffer(self):
+    buf = b'\x01\x02\x03'
+    expected = np.frombuffer(buf, dtype='uint8')
+    actual = jnp.frombuffer(buf, dtype='uint8')
+    self.assertArraysEqual(expected, actual)
+
+  def testFromFunction(self):
+    def f(x, y, z):
+      return x + 2 * y + 3 * z
+    shape = (3, 4, 5)
+    expected = np.fromfunction(f, shape=shape)
+    actual = jnp.fromfunction(f, shape=shape)
+    self.assertArraysEqual(expected, actual)
+
+  def testFromString(self):
+    s = "1,2,3"
+    expected = np.fromstring(s, sep=',', dtype=int)
+    actual = jnp.fromstring(s, sep=',', dtype=int)
+    self.assertArraysEqual(expected, actual)
+
+
 # Most grad tests are at the lax level (see lax_test.py), but we add some here
 # as needed for e.g. particular compound ops of interest.
 
@@ -6093,6 +6114,7 @@ def testWrappedSignaturesMatch(self):
       'identity': ['like'],
       'full': ['order', 'like'],
       'full_like': ['subok', 'order'],
+      'fromfunction': ['like'],
       'histogram': ['normed'],
       'histogram2d': ['normed'],
       'histogramdd': ['normed'],