internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: google/jax

Issue:
forward-mode eigh derivatives return NaN's
First off, thanks for developing JAX, it's awesome!

I recently noticed that when one computes higher order derivatives of expressions involving eigenvectors obtained from `jax.numpy.eigh` or `jax.scipy.eigh`, NaNs are returned. However, it only appears to be a problem with forward-mode autodiff. This is unfortunate, since forward-mode tends to be more efficient for higher order derivatives.

A minimal example:
```python
import jax
import jax.numpy as np

# Create a random symmetric matrix
seed = jax.random.PRNGKey(0)
tmp = jax.random.uniform(seed, (2,2))
a = np.dot(tmp, tmp.T)

def test(inp):
    val, vec = np.linalg.eigh(inp)
    return np.dot(np.dot(vec, inp), vec.T)

def test_deriv(func, inp):
    grad_func = jax.jacfwd(func)
    hess_func = jax.jacfwd(grad_func)
    cube_func = jax.jacfwd(hess_func) # This derivative returns NaN, but jax.jacrev works!
    print(grad_func(inp))
    print(hess_func(inp))
    print(cube_func(inp))

test_deriv(test, a)
```

Additional notes:
- Expressions involving eigenvalues do not appear to have this issue
- Any array contractions involving eigenvectors and other arrays (`dot`, `einsum`, etc) also exhibit similar behavior: `jacfwd` fails, while `jacrev` works.
- I'm using CPU
- JAX version 0.1.46

Thanks for your time!


Discussion:
The first NaN in the computation appears to come from the forward-mode derivative of `pow` where the LHS is negative:
`jax.jvp(lax.pow, (-1.5, 2.), (0., 0.))`

This in turn comes from the `y ** 2` term in the RHS JVP for `div`.



Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1720,7 +1720,7 @@ def _div_transpose_rule(cotangent, x, y):
 div_p = standard_binop([_num, _num], 'div')
 ad.defjvp(div_p,
           lambda g, x, y: div(_brcast(g, y), y),
-          lambda g, x, y: div(mul(neg(_brcast(g, x)), x), pow(y, _two(y))))
+          lambda g, x, y: div(mul(neg(_brcast(g, x)), x), square(y)))
 ad.primitive_transposes[div_p] = _div_transpose_rule
 
 rem_p = standard_binop([_num, _num], 'rem')


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -538,6 +538,20 @@ def testIssue1151(self):
     jac0 = jax.jacobian(np.linalg.solve, argnums=0)(A[0], b[0])
     jac1 = jax.jacobian(np.linalg.solve, argnums=1)(A[0], b[0])
 
+  def testIssue1383(self):
+    seed = jax.random.PRNGKey(0)
+    tmp = jax.random.uniform(seed, (2,2))
+    a = np.dot(tmp, tmp.T)
+
+    def f(inp):
+      val, vec = np.linalg.eigh(inp)
+      return np.dot(np.dot(vec, inp), vec.T)
+
+    grad_func = jax.jacfwd(f)
+    hess_func = jax.jacfwd(grad_func)
+    cube_func = jax.jacfwd(hess_func)
+    self.assertFalse(onp.any(onp.isnan(cube_func(a))))
+
 
 class ScipyLinalgTest(jtu.JaxTestCase):
Repository: google/jax

Issue:
advi.py example is broken with jit
advi.py example is broken with jit: @partial(jit, static_argnums=(0, 1, 2, 4))
If jit is removed, then It works.


Discussion:
This is because the `4`th positional argument, `num_ticks`, has a default value and not explicitly supplied during the invocation. 

Just setting the `static_argnums` to (0,1,2) would fix this. 

Resolving patch:
diff --git a/examples/advi.py b/examples/advi.py
--- a/examples/advi.py
+++ b/examples/advi.py
@@ -54,7 +54,7 @@ def batch_elbo(logprob, rng, params, num_samples):
 # ========= Helper function for plotting. =========
 
 @partial(jit, static_argnums=(0, 1, 2, 4))
-def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):
+def _mesh_eval(func, x_limits, y_limits, params, num_ticks):
     # Evaluate func on a 2D grid defined by x_limits and y_limits.
     x = np.linspace(*x_limits, num=num_ticks)
     y = np.linspace(*y_limits, num=num_ticks)
@@ -63,6 +63,8 @@ def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):
     zs = vmap(func, in_axes=(0, None))(xy_vec, params)
     return X, Y, zs.reshape(X.shape)
 
+def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):
+    return _mesh_eval(func, x_limits, y_limits, params, num_ticks)
 
 # ========= Define an intractable unnormalized density =========
 


Test patch:
Repository: google/jax

Issue:
psum of complex64 not supported for multi GPU.
Complex numbers are not supported for AllReduce.

I found this bug when trying to do psum across 4 GPUs with complex64 as the dtype.

```
RuntimeError: Unimplemented: Requested AllReduce not implemented on GPU; replica_count: 4; operand_count: 1; IsCrossReplicaAllReduce: 1; NCCL support: 1; first operand array element-type: C64
```


Discussion:
We can fix this at the JAX level by making [the `psum` lowering](https://github.com/google/jax/blob/01f81875623d7548402d0224703a70e9052492b5/jax/lax/lax_parallel.py#L190-L194) switch on the array element type, and just generate a pair of AllReduces for complex types.

Resolving patch:
diff --git a/jax/lax/lax_parallel.py b/jax/lax/lax_parallel.py
--- a/jax/lax/lax_parallel.py
+++ b/jax/lax/lax_parallel.py
@@ -193,11 +193,20 @@ def _allreduce_translation_rule(prim, c, val, replica_groups, backend=None):
   computation = xla.primitive_computation(prim, scalar, scalar, backend=backend)
   return c.AllReduce(val, computation, replica_groups=replica_groups)
 
+# psum translation rule has special handling for complex dtypes
+def _psum_translation_rule(c, val, replica_groups, backend=None):
+  psum = partial(_allreduce_translation_rule, lax.add_p, c,
+                 replica_groups=replica_groups, backend=backend)
+  dtype = c.GetShape(val).numpy_dtype()
+  if onp.issubdtype(dtype, onp.complexfloating):
+    return c.Complex(psum(c.Real(val)), psum(c.Imag(val)))
+  else:
+    return psum(val)
+
 psum_p = standard_pmap_primitive('psum')
 pxla.split_axis_rules[psum_p] = \
     partial(_allreduce_split_axis_rule, psum_p, lax._reduce_sum)
-xla.parallel_translations[psum_p] = \
-    partial(_allreduce_translation_rule, lax.add_p)
+xla.parallel_translations[psum_p] = _psum_translation_rule
 pxla.parallel_pure_rules[psum_p] = lambda x, shape: x * prod(shape)
 ad.deflinear(psum_p, lambda t, axis_name: [t])
 


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -66,6 +66,17 @@ def testBasic(self):
     ans = f(x)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testComplexPsum(self):
+    f = pmap(lambda x: x - lax.psum(x, 'i'), axis_name='i')
+
+    shape = (xla_bridge.device_count(), 4 * 2)
+    x = onp.arange(prod(shape), dtype=onp.float32).reshape(shape).view(onp.complex64)
+    expected = x - onp.sum(x, 0)
+
+    ans = f(x)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+
   def testNestedBasic(self):
     f = lambda x: lax.psum(lax.psum(x, 'i'), 'j')
     f = pmap(pmap(f, 'i'), 'j')
Repository: google/jax

Issue:
Document pmap on pytrees
The pmap documentation does not mention how it can be used on pytrees or what it is supposed to do when pytrees are involved. For example, code like what is pasted below in Trax seems to do things on pytrees.

``` 
@functools.partial(backend.pmap, axis_name='batch')
def mapped_predict(x, params, state, rng):
    return model_predict(x, params=params, state=state, rng=rng)
```




Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -652,7 +652,7 @@ def pmap(fun, axis_name=None, devices=None, backend=None):
   devices available, as returned by ``jax.local_device_count()`` (unless
   ``devices`` is specified, see below). For nested ``pmap`` calls, the product
   of the mapped axis sizes must be less than or equal to the number of XLA
-  devices.  TODO(skye): support < # local devices on multi-host platforms
+  devices.
 
   **Multi-host platforms:** On multi-host platforms such as TPU pods, ``pmap``
   is designed to be used in SPMD Python programs, where every host is running
@@ -667,7 +667,9 @@ def pmap(fun, axis_name=None, devices=None, backend=None):
   "sees" only its local shard of the input and output.
 
   Args:
-    fun: Function to be mapped over argument axes.
+    fun: Function to be mapped over argument axes. Its arguments and return
+      value should be arrays, scalars, or (nested) standard Python containers
+      (tuple/list/dict) thereof.
     axis_name: Optional, a hashable Python object used to identify the mapped
       axis so that parallel collectives can be applied.
     devices: This is an experimental feature and the API is likely to change.


Test patch:
Repository: google/jax

Issue:
np.broadcast_to regression in the latest JAX version
The new version 0.1.47 introduces a regression where `np.broadcast_to(1, (3, 2))` throws the error `AttributeError: 'int' object has no attribute 'ndim'`.

cc @mattjj @neerajprad


Discussion:

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -599,7 +599,7 @@ def broadcast(operand, sizes):
   return broadcast_p.bind(operand, sizes=tuple(sizes))
 
 def broadcast_in_dim(operand, shape, broadcast_dimensions):
-  if operand.ndim == len(shape) and not len(broadcast_dimensions):
+  if onp.ndim(operand) == len(shape) and not len(broadcast_dimensions):
     return operand
   if any(x < 0 or x >= len(shape) for x in broadcast_dimensions):
     msg = ("broadcast dimensions must be >= 0 and < ndim(shape), got {} for "


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
@@ -1952,6 +1952,10 @@ def testBroadcastToIssue1522(self):
         ValueError, "Incompatible shapes for broadcasting: .*",
         lambda: lnp.broadcast_to(onp.ones((2, 3)), (1, 3)))
 
+  def testBroadcastToIntIssue1548(self):
+    self.assertAllClose(lnp.broadcast_to(1, (3, 2)), onp.ones((3, 2)),
+                        check_dtypes=False)
+
 # 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.
Repository: google/jax

Issue:
don't generate trivial gathers from numpy indexing
A NumPy expression like `x[:, None]` should only produce a `lax.reshape` call, but currently it produces a trivial `lax.gather` followed by the `lax.reshape`. We should avoid generating the trivial gather.

We'll need to change lax_numpy.py's `_gather` or `_index_to_gather` to detect when a gather is trivial.


Discussion:
Maybe we just need this `if` check:

```python
  if indexer.gather_indices.size:
    y = lax.gather(y, indexer.gather_indices, indexer.dnums,
                   indexer.gather_slice_shape)
```

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -2548,9 +2548,11 @@ def _rewriting_take(arr, idx):
 def _gather(arr, treedef, static_idx, dynamic_idx):
   idx = _merge_static_and_dynamic_indices(treedef, static_idx, dynamic_idx)
   indexer = _index_to_gather(shape(arr), idx)  # shared with _scatter_update
+  y = arr
 
-  y = lax.gather(arr, indexer.gather_indices, indexer.dnums,
-                 indexer.gather_slice_shape)
+  if indexer.gather_indices.size:
+    y = lax.gather(y, indexer.gather_indices, indexer.dnums,
+                   indexer.gather_slice_shape)
 
   # Reverses axes with negative strides.
   if indexer.reversed_y_dims:


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -730,6 +730,11 @@ def testJVPOfGradOfIndexing(self):
     self.assertAllClose(expected, primals, check_dtypes=True)
     self.assertAllClose(onp.zeros_like(x), tangents, check_dtypes=True)
 
+  def testTrivialGatherIsntGenerated(self):
+    # https://github.com/google/jax/issues/1621
+    jaxpr = api.make_jaxpr(lambda x: x[:, None])(onp.arange(4))
+    self.assertEqual(len(jaxpr.eqns), 1)
+    self.assertNotIn('gather', str(jaxpr))
 
 
 def _broadcastable_shapes(shape):
@@ -891,6 +896,5 @@ def testSegmentSum(self):
     self.assertAllClose(ans, expected, check_dtypes=False)
 
 
-
 if __name__ == "__main__":
   absltest.main()
Repository: google/jax

Issue:
ppermute is returning 0s
Tested with TPUs.

```python
@partial(jax.pmap, axis_name="shard")
def f(a):
  return jax.lax.ppermute(
      a, "shard", zip(range(8), [1, 0, 2, 3, 4, 5, 6, 7]))

a = np.ones((8, 2, 2))
print(f(a))
```

Prints out
```
[[[0. 0.]
  [0. 0.]]

 [[0. 0.]
  [0. 0.]]

 [[0. 0.]
  [0. 0.]]

 [[0. 0.]
  [0. 0.]]

 [[0. 0.]
  [0. 0.]]

 [[0. 0.]
  [0. 0.]]

 [[0. 0.]
  [0. 0.]]

 [[0. 0.]
  [0. 0.]]]
```


Discussion:
@hawkinsp pointed out that adding `tuple(...)` around the `zip(...)` fixes it. For a quick fix, we can always explicitly cast the `perm` argument to `ppermute` to a `tuple`. I'll make a PR for this.

Resolving patch:
diff --git a/jax/lax/lax_parallel.py b/jax/lax/lax_parallel.py
--- a/jax/lax/lax_parallel.py
+++ b/jax/lax/lax_parallel.py
@@ -114,7 +114,7 @@ def ppermute(x, axis_name, perm):
     An array with the same shape as ``x`` with slices along the axis
     ``axis_name`` gathered from ``x`` according to the permutation ``perm``.
   """
-  return ppermute_p.bind(x, axis_name=axis_name, perm=perm)
+  return ppermute_p.bind(x, axis_name=axis_name, perm=tuple(perm))
 
 def pswapaxes(x, axis_name, axis):
   """Swap the pmapped axis ``axis_name`` with the unmapped axis ``axis``.


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -331,6 +331,16 @@ def testCollectivePermuteCyclicGrad(self):
     expected = onp.roll(onp.pi + onp.arange(device_count), 1)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  @jtu.skip_on_devices("cpu", "gpu")
+  def testIssue1703(self):
+    num_devices = xla_bridge.device_count()
+    perm = [num_devices - 1] + list(range(num_devices - 1))
+    f = pmap(
+      lambda x: lax.ppermute(x, "i", zip(range(num_devices), perm)), "i")
+    result = f(np.arange(num_devices, dtype=np.float32))
+    expected = np.asarray(perm, dtype=np.float32)
+    self.assertAllClose(result, expected)
+
   @jtu.skip_on_devices("cpu", "gpu")
   def testRule30(self):
     # This is a test of collective_permute implementing a simple halo exchange
Repository: google/jax

Issue:
unary + does not work in jax
```
>>> import jax
>>> +jax.numpy.asarray([1,2,3])
/Users/lukasheinrich/Code/excursiondev/excvenv/lib/python3.6/site-packages/jax/lib/xla_bridge.py:119: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: bad operand type for unary +: 'DeviceArray'
```

while in numpy it works

```
>>> import numpy
>>> +numpy.asarray([1,2,3])
array([1, 2, 3])
```

a similar issue appear for the `-`  operator


Discussion:
I see the issue for unary `+`; that's just a case we forgot to implement (the `+` operator isn't that useful!) Do you have an example of unary `-` failing? It works for me.
sorry, I misremembered: unary minus works for me.

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -300,6 +300,7 @@ def aval(self):
     assert False
 
   def __neg__(self): return self.aval._neg(self)
+  def __pos__(self): return self.aval._pos(self)
   def __eq__(self, other): return self.aval._eq(self, other)
   def __ne__(self, other): return self.aval._ne(self, other)
   def __lt__(self, other): return self.aval._lt(self, other)
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -3194,6 +3194,7 @@ def _unimplemented_setitem(self, i, x):
     "getitem": _rewriting_take,
     "setitem": _unimplemented_setitem,
     "neg": negative,
+    "pos": positive,
     "eq": equal,
     "ne": not_equal,
     "lt": less,


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
@@ -256,6 +256,7 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
     op_record("__lt__", 2, default_dtypes, all_shapes, jtu.rand_default, []),
     op_record("__gt__", 2, default_dtypes, all_shapes, jtu.rand_default, []),
     op_record("__ge__", 2, default_dtypes, all_shapes, jtu.rand_default, []),
+    op_record("__pos__", 1, number_dtypes, all_shapes, jtu.rand_default, []),
     op_record("__neg__", 1, number_dtypes, all_shapes, jtu.rand_default, []),
     op_record("__pow__", 2, inexact_dtypes, all_shapes, jtu.rand_positive, [],
               tolerance={onp.float32: 2e-4, onp.complex64: 2e-4, onp.complex128: 1e-14}),
Repository: google/jax

Issue:
ImportError: ..... ____chkstk_darwin
Hi folks,
Wondering if someone can please help with this import error on a fresh install. 

Background: Recently started a new conda environment with a couple of basics like jupyter, numpy, scipy. I'm using MacOS High Sierra 10.13.4, and python3.7

The specific error is:
```
---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
<ipython-input-1-77da20ac745a> in <module>
----> 1 from jax import vmap # for auto-vectorizing functions
      2 from functools import partial # for use with vmap
      3 from jax import jit # for compiling functions for speedup
      4 from jax import random # stax initialization uses jax.random
      5 from jax.experimental import stax # neural network library
```
...

```
~/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_client.py in <module>
     36 # and TensorFlow may fail with duplicate protocol buffer message definitions.
     37 
---> 38 from . import xla_extension as _xla
     39 from .xla_extension import ops
     40 

ImportError: dlopen(/Users/lmar3213/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_extension.so, 2): Symbol not found: ____chkstk_darwin
  Referenced from: /Users/lmar3213/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_extension.so (which was built for Mac OS X 10.15)
  Expected in: /usr/lib/libSystem.B.dylib
 in /Users/lmar3213/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_extension.so
```

To install, I ran as per the installation instructions i.e. `pip install --upgrade pip` which was already up to date, and then `pip install --upgrade jax jaxlib` which installed happily. 

The only clue I have as to understanding the error is that it mentions `xla_extension.so` was built for mac OS X 10.15 but Im on 10.13. Any help is appreciated! Thankyou!



Discussion:

Resolving patch:
diff --git a/build/build.py b/build/build.py
--- a/build/build.py
+++ b/build/build.py
@@ -176,6 +176,7 @@ def check_bazel_version(bazel_path, min_version, max_version):
 
 # Sets the default Apple platform to macOS.
 build --apple_platform_type=macos
+build --macos_minimum_os=10.9
 
 # Make Bazel print out all options from rc files.
 build --announce_rc


Test patch:
Repository: google/jax

Issue:
Indexing into an empty array fails
Indexing a non-empty array with an empty array of indices works fine:

```
import jax.numpy as np
values = np.zeros(5, dtype=np.float32)
indices = np.array([], dtype=np.int32)
values[indices]

-> DeviceArray([], dtype=float32)
```

But trying to index an empty array with an empty array of indices causes a runtime error prompting me to file an issue:

```
import jax.numpy as np
values = np.zeros(0, dtype=np.float32)
indices = np.array([], dtype=np.int32)
values[indices]

...

RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 1), got 1.: 
This is a bug in JAX's shape-checking rules; please report it!
https://github.com/google/jax/issues
```

The same code works fine in vanilla numpy:

```
import numpy as np
values = np.zeros(0, dtype=np.float32)
indices = np.array([], dtype=np.int32)
values[indices]
-> array([], dtype=float32)
```

This is on jax 0.1.49.


Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -2656,9 +2656,13 @@ def _gather(arr, treedef, static_idx, dynamic_idx):
   indexer = _index_to_gather(shape(arr), idx)  # shared with _scatter_update
   y = arr
 
-  # We avoid generating a gather when indexer.gather_indices.size is empty
-  # unless indexer.slice_shape also corresponds to an empty array.
-  if indexer.gather_indices.size or not _prod(indexer.slice_shape):
+  # Avoid calling gather if the slice shape is empty, both as a fast path and to
+  # handle cases like zeros(0)[array([], int32)].
+  if _prod(indexer.slice_shape) == 0:
+    return zeros(indexer.slice_shape, dtype=y.dtype)
+
+  # We avoid generating a gather when indexer.gather_indices.size is empty.
+  if indexer.gather_indices.size:
     y = lax.gather(y, indexer.gather_indices, indexer.dnums,
                    indexer.gather_slice_shape)
 


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -223,6 +223,7 @@ def check_grads(f, args, order, atol=None, rtol=None, eps=None):
      IndexSpec(shape=(3, 4, 5), indexer=onp.array([0, 2, 0, 1])),
      IndexSpec(shape=(3,), indexer=onp.array([-1, 1])),
      IndexSpec(shape=(3,), indexer=onp.array([-2, -1])),
+     IndexSpec(shape=(0,), indexer=onp.array([], dtype=onp.int32)),
      ]),
     ("One2DIntArrayIndex",
      [IndexSpec(shape=(3,), indexer=onp.array([[0, 0]])),
@@ -274,6 +275,7 @@ def check_grads(f, args, order, atol=None, rtol=None, eps=None):
       IndexSpec(shape=(3, 4, 5), indexer=onp.array([0, 2, 1])),
       IndexSpec(shape=(3,), indexer=onp.array([-1, 1])),
       IndexSpec(shape=(3,), indexer=onp.array([-2, -1])),
+      IndexSpec(shape=(0,), indexer=onp.array([], dtype=onp.int32)),
      ]),
     ("One2DIntArrayIndex",
      [IndexSpec(shape=(3,), indexer=onp.array([[0, 1]])),
Repository: google/jax

Issue:
not compatible with numpy 1.18 and python 3.7
With numpy 1.18 and python 3.7, `import jax` will raise the error
```python
>>> import jax
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/fehiepsi/miniconda3/envs/py37/lib/python3.7/site-packages/jax/__init__.py", line 20, in <module>
    from jax import nn
  File "/home/fehiepsi/miniconda3/envs/py37/lib/python3.7/site-packages/jax/nn/__init__.py", line 17, in <module>
    from . import initializers
  File "/home/fehiepsi/miniconda3/envs/py37/lib/python3.7/site-packages/jax/nn/initializers.py", line 28, in <module>
    from jax import random
  File "/home/fehiepsi/miniconda3/envs/py37/lib/python3.7/site-packages/jax/random.py", line 33, in <module>
    from . import numpy as np
  File "/home/fehiepsi/miniconda3/envs/py37/lib/python3.7/site-packages/jax/numpy/__init__.py", line 16, in <module>
    from .lax_numpy import *
  File "<frozen importlib._bootstrap>", line 1019, in _handle_fromlist
  File "/home/fehiepsi/miniconda3/envs/py37/lib/python3.7/site-packages/jax/numpy/lax_numpy.py", line 3211, in wrapped
    raise NotImplementedError(msg.format(fun))
NotImplementedError: Numpy function <function __getattr__ at 0x7f029e9dd560> not yet implemented
```


Discussion:

Resolving patch:
diff --git a/jax/util.py b/jax/util.py
--- a/jax/util.py
+++ b/jax/util.py
@@ -214,6 +214,10 @@ def get_module_functions(module):
   """
   module_fns = set()
   for key in dir(module):
+    # Omitting module level __getattr__, __dir__ which was added in Python 3.7
+    # https://www.python.org/dev/peps/pep-0562/
+    if key in ('__getattr__', '__dir__'):
+      continue
     attr = getattr(module, key)
     if isinstance(
         attr, (types.BuiltinFunctionType, types.FunctionType, onp.ufunc)):


Test patch: