Repository: jax-ml/jax
Pull request: Make JAX compatible with numpy version 1.23.0
For more discussion of these changes, see #11221

Fixes #11221

Problem statement:
autodiff fails under numpy 1.23
Something in the numpy 1.23 release broke autodiff; in particular, float0 arrays are being converted to bool:
```python
import jax
import numpy
print(numpy.__version__)
print(jax.grad(lambda x: x+0., allow_int=True)(1).dtype)
```
Newest numpy:
```
1.23.0
bool
```
previous version:
```
1.22.4
[('float0', 'V')]
```

Hints from issue discussion:
Temporary pinning to fix CI in https://github.com/google/jax/pull/11220
Located the issue to this line: https://github.com/google/jax/blob/b623ed58b0415f31b4c2f2c389220f1e51de8d93/jax/interpreters/ad.py#L65

In numpy 1.22, `dtype(t) is float0` is True, and in numpy 1.23 `dtype(t) is float0` is False.
This is the operative change in numpy:
```python
import numpy as np
from jax import dtypes
print(np.__version__)
print(np.result_type(dtypes.float0) is dtypes.float0)
```
Before:
```
1.22.4
True
```
After:
```
1.23.0
False
```
I don't know for sure, but I suspect this is the relevant change: https://github.com/numpy/numpy/pull/19345/

Edit: actually, the relevant change is almost certainly this: https://github.com/numpy/numpy/pull/19226 ("API: Fix structured dtype cast-safety, promotion, and comparison")
OK, I think this is the result of the new `ensure_canonical` functionality added in https://github.com/numpy/numpy/pull/19226, along with https://github.com/numpy/numpy/pull/19346 which implements `ensure_canonical` for structured arrays.

Part of the reason I think this is because of this little quirk:
```python
In [1]: from jax.dtypes import float0                                                                                                                                                                                                                                                                            

In [2]: import numpy as np                                                                                                                                                                                                                                                                                       

In [3]: np.result_type(float0) is float0                                                                                                                                                                                                                                                                         
Out[3]: False

In [4]: np.result_type(float0, float0) is float0                                                                                                                                                                                                                                                                 
Out[4]: True
```
Along with the fact that `NPY_DT_CALL_ensure_canonical` is only called on the result for the single-input case: https://github.com/numpy/numpy/blob/89b942e41b45ddd91762e704d3f2ca5677a320b7/numpy/core/src/multiarray/convert_datatype.c#L1633

I'm not sure exactly the mechanism whereby this generates a copy, but it seems clear that the fix could go one of two ways:

1. Somehow change our definition of `float0` so it is in a canonical form that doesn't generate a copy when `NPY_DT_CALL_ensure_canonical` is called, or
2. Stop relying on `is` identity to determine if a dtype is equal to `float0`

Option 2 is probably safer in the long run, as I can find no explicit guarantee that functions like `np.dtype` and `np.result_type` will return no-copy references for non-built-in dtypes. However, equality checks for numpy dtypes are annoyingly polymorphic, so this must be used with care (probably `dtype is float0` should be replaced with `np.dtype(dtype) == float0` for safety). Also, since equality check is slightly more expensive than identity check, and since this is used so often throughout the JAX package, this change may have a trace-time performance impact.

I'm going to look deeper and try to find exactly where the copy is being generated, but unless something else comes to light, I think option 2 is our best bet.
I believe the `ensure_canonical` check for void types leads to this function: https://github.com/numpy/numpy/blob/d4baf3641471971552206b02f94323a54b54d99a/numpy/core/src/multiarray/dtypemeta.c#L295 which ends up generating a copy for JAX's definition of `float0`. This was changed a few months back in https://github.com/numpy/numpy/pull/19346

Here's the TODO that broke our use-case: https://github.com/numpy/numpy/blame/d4baf3641471971552206b02f94323a54b54d99a/numpy/core/src/multiarray/dtypemeta.c#L324-L325

I don't see any way for us to prevent copy on dtype canonicalization, unless we define `float0` as something other than a structured void dtype.
Additionally, there are at least two regressions in numpy v1.23 that are causing test failures in the JAX test suite: https://github.com/numpy/numpy/issues/21838 and https://github.com/numpy/numpy/issues/21840

Commit message: TST: skip tests on numpy 1.23.0 due to regressions in that release

Patch to review:
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index 43a46d68d068..6af5643f1a5c 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -957,7 +957,7 @@ def _maybe_perturbed(x: Any) -> bool:
     # happen later, but some types always have trivial tangents.
     vspace = x.aval.at_least_vspace()
     return not (vspace is core.abstract_token or
-                getattr(vspace, 'dtype', None) is dtypes.float0)
+                getattr(vspace, 'dtype', None) == dtypes.float0)
   elif not isinstance(x, ad.JVPTracer):
     # If x is not a JVPTracer, recursively check its contents.
     return any(_maybe_perturbed(attr) for name, attr in x._contents())
diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
index 27997eea8069..9dcb0c45acd7 100644
--- a/jax/_src/dispatch.py
+++ b/jax/_src/dispatch.py
@@ -620,7 +620,7 @@ def aval_to_result_handler(sticky_device: Optional[Device],
 
 def array_result_handler(sticky_device: Optional[Device],
                          aval: core.ShapedArray):
-  if aval.dtype is dtypes.float0:
+  if aval.dtype == dtypes.float0:
     return lambda _, __: np.zeros(aval.shape, dtypes.float0)
   aval = core.raise_to_shaped(aval)
   handler = lambda _, b: device_array.make_device_array(aval, sticky_device, b)
@@ -629,7 +629,7 @@ def array_result_handler(sticky_device: Optional[Device],
 
 def dynamic_array_result_handler(sticky_device: Optional[Device],
                                  aval: core.DShapedArray):
-  if aval.dtype is dtypes.float0:
+  if aval.dtype == dtypes.float0:
     return lambda _: np.zeros(aval.shape, dtypes.float0)  # type: ignore
   else:
     return partial(_dynamic_array_result_handler, sticky_device, aval)
@@ -974,7 +974,7 @@ def device_put(x, device: Optional[Device] = None) -> Tuple[Any]:
 
 def _device_put_array(x, device: Optional[Device]):
   backend = xb.get_device_backend(device)
-  if x.dtype is dtypes.float0:
+  if x.dtype == dtypes.float0:
     x = np.zeros(x.shape, dtype=np.dtype(bool))
   return (backend.buffer_from_pyval(x, device),)
 
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 6f68c06aeef3..154dfb0188a0 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -1405,7 +1405,7 @@ def naryop_dtype_rule(result_dtype, accepted_dtypes, name, *avals, **kwargs):
   aval_dtypes = [aval.dtype for aval in avals]
   for i, (aval_dtype, types) in enumerate(zip(aval_dtypes, accepted_dtypes)):
     if not any(dtypes.issubdtype(aval_dtype, t) for t in types):
-      if aval_dtype is dtypes.float0:
+      if aval_dtype == dtypes.float0:
         raise TypeError(
             f"Called {name} with a float0 at position {i}. "
             "float0s do not support any operations by design, because they "
@@ -1446,8 +1446,8 @@ def _broadcasting_shape_rule(name, *avals):
   return tuple(result_shape)
 
 def _naryop_weak_type_rule(name, *avals, **kwargs):
-  if any(aval.dtype is dtypes.float0 for aval in avals):
-    pos = next(i for i, aval in enumerate(avals) if aval.dtype is dtypes.float0)
+  if any(aval.dtype == dtypes.float0 for aval in avals):
+    pos = next(i for i, aval in enumerate(avals) if aval.dtype == dtypes.float0)
     raise TypeError(
         f"Called {name} with a float0 at position {pos}. "
         "float0s do not support any operations by design, because they "
@@ -2206,14 +2206,14 @@ def _convert_element_type_transpose_rule(ct, operand, *, new_dtype, weak_type):
   old_weak_type = dtypes.is_weakly_typed(operand)
   if type(ct) is ad_util.Zero:
     return [ad_util.Zero(operand.aval)]
-  elif core.primal_dtype_to_tangent_dtype(old_dtype) is dtypes.float0:
+  elif core.primal_dtype_to_tangent_dtype(old_dtype) == dtypes.float0:
     return [ad_util.Zero(operand.aval.update(dtype=dtypes.float0, weak_type=False))]
   else:
     return [convert_element_type_p.bind(ct, new_dtype=old_dtype,
                                         weak_type=old_weak_type)]
 
 def _convert_element_type_jvp_rule(tangent, operand , *, new_dtype, weak_type):
-  if core.primal_dtype_to_tangent_dtype(new_dtype) is dtypes.float0:
+  if core.primal_dtype_to_tangent_dtype(new_dtype) == dtypes.float0:
     return ad_util.Zero(tangent.aval.update(dtype=dtypes.float0, weak_type=False))
   else:
     return convert_element_type_p.bind(tangent, new_dtype=new_dtype,
diff --git a/jax/_src/numpy/util.py b/jax/_src/numpy/util.py
index af88ffbdaf8d..be231b13d792 100644
--- a/jax/_src/numpy/util.py
+++ b/jax/_src/numpy/util.py
@@ -314,7 +314,7 @@ def _check_arraylike(fun_name, *args):
 
 def _check_no_float0s(fun_name, *args):
   """Check if none of the args have dtype float0."""
-  if any(dtypes.dtype(arg) is dtypes.float0 for arg in args):
+  if any(dtypes.dtype(arg) == dtypes.float0 for arg in args):
     raise TypeError(
         f"Called {fun_name} with a float0 array. "
         "float0s do not support any operations by design because they "
diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py
index 84599aa1a740..3ae212c31bd1 100644
--- a/jax/_src/test_util.py
+++ b/jax/_src/test_util.py
@@ -304,6 +304,7 @@ def format_test_name_suffix(opname, shapes, dtypes):
 # between NumPy scalars, Python scalars, and 0-D arrays.
 class ScalarShape:
   def __len__(self): return 0
+  def __getitem__(self, i): raise IndexError(f"index {i} out of range.")
 class _NumpyScalar(ScalarShape): pass
 class _PythonScalar(ScalarShape): pass
 NUMPY_SCALAR_SHAPE = _NumpyScalar()
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
index 90c08067d418..9f2c23657bff 100644
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -65,7 +65,7 @@ def jvp(fun: lu.WrappedFun, has_aux=False, instantiate=True,
 @lu.transformation
 def jvpfun(instantiate, transform_stack, primals, tangents):
   tangents = [Zero.from_value(t) if not isinstance(t, Zero)
-              and dtype(t) is float0 else t for t in tangents]
+              and dtype(t) == float0 else t for t in tangents]
   ctx = (source_info_util.transform_name_stack('jvp') if transform_stack
          else contextlib.nullcontext())
   with core.new_main(JVPTrace) as main, ctx:
@@ -158,7 +158,7 @@ def unpair_pval(pval):
     return (aval_1, const_1), (aval_2, const_2)
 
 def replace_float0s(primal, tangent):
-  if dtype(tangent) is float0:
+  if dtype(tangent) == float0:
     return zeros_like_jaxval(primal)
   else:
     return tangent
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
index 140027674b02..55e8842ca08e 100644
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -80,7 +80,7 @@ def identity(x): return x
 _scalar_types = dtypes.python_scalar_dtypes.keys()
 
 def _make_array_shape(a: ShapedArray) -> Sequence[XlaShape]:
-  if a.dtype is dtypes.float0:
+  if a.dtype == dtypes.float0:
     return (xc.Shape.array_shape(np.dtype('bool'), a.shape),)
   else:
     return (xc.Shape.array_shape(a.dtype, a.shape),)
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 430f2cde6b89..a8a7a11b2bd9 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -104,7 +104,7 @@ def _shape_and_dtypes(shapes, dtypes):
       yield (shape, dtype)
 
 def _compatible_shapes(shape):
-  if shape in scalar_shapes or np.ndim(shape) == 0:
+  if np.ndim(shape) == 0 or shape in scalar_shapes:
     return [shape]
   return (shape[n:] for n in range(len(shape) + 1))
 
@@ -2395,6 +2395,9 @@ def testDeleteMaskArray(self, shape, dtype, axis):
     rng = jtu.rand_default(self.rng())
     mask_size = np.zeros(shape).size if axis is None else np.zeros(shape).shape[axis]
     mask = jtu.rand_int(self.rng(), low=0, high=2)(mask_size, bool)
+    if numpy_version == (1, 23, 0) and mask.shape == (1,):
+      # https://github.com/numpy/numpy/issues/21840
+      self.skipTest("test fails for numpy v1.23.0")
     args_maker = lambda: [rng(shape, dtype)]
     np_fun = lambda arg: np.delete(arg, mask, axis=axis)
     jnp_fun = lambda arg: jnp.delete(arg, mask, axis=axis)
@@ -2622,6 +2625,9 @@ def np_fun(x, fill_value=fill_value):
       {"testcase_name": f"_{dtype.__name__}", "dtype": dtype}
       for dtype in inexact_dtypes))
   def testUniqueNans(self, dtype):
+    if numpy_version == (1, 23, 0) and dtype == np.float16:
+      # https://github.com/numpy/numpy/issues/21838
+      self.skipTest("Known failure on numpy 1.23.0")
     def args_maker():
       x = [-0.0, 0.0, 1.0, 1.0, np.nan, -np.nan]
       if np.issubdtype(dtype, np.complexfloating):
@@ -6333,6 +6339,7 @@ def testWrappedSignaturesMatch(self):
     # TODO(jakevdp): fix some of the following signatures. Some are due to wrong argument names.
     unsupported_params = {
       'asarray': ['like'],
+      'average': ['keepdims'],
       'broadcast_to': ['subok', 'array'],
       'clip': ['kwargs'],
       'copy': ['subok'],
@@ -6352,6 +6359,7 @@ def testWrappedSignaturesMatch(self):
       'ones': ['order', 'like'],
       'ones_like': ['subok', 'order'],
       'tri': ['like'],
+      'unique': ['equal_nan'],
       'zeros_like': ['subok', 'order']
     }
 
diff --git a/tests/random_test.py b/tests/random_test.py
index 07e418c397d7..9840e8625631 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -643,7 +643,7 @@ def testChoice(self, dtype, input_range_or_shape, shape, replace, weighted, axis
     key = self.seed_prng(0)
     is_range = type(input_range_or_shape) is int
     x = (input_range_or_shape if is_range else
-         self.rng().permutation(jnp.arange(np.prod(
+         self.rng().permutation(np.arange(np.prod(
            input_range_or_shape), dtype=dtype)).reshape(input_range_or_shape))
     N = x if is_range else x.shape[axis]
     p = None if not weighted else (np.arange(N) + 1) / np.sum(np.arange(N) + 1)
@@ -678,7 +678,7 @@ def testPermutation(self, dtype, range_or_shape, axis, independent):
     key = self.seed_prng(0)
     is_range = type(range_or_shape) is int
     x = (range_or_shape if is_range else
-         self.rng().permutation(jnp.arange(
+         self.rng().permutation(np.arange(
            np.prod(range_or_shape), dtype=dtype)).reshape(range_or_shape))
     shape = ((range_or_shape,) if is_range else range_or_shape)
     x_ = np.copy(x)


Code review comments:
diff_hunk: @@ -6,4 +6,4 @@ pillow>=8.3.1,<9.1.0
 pytest-benchmark
 pytest-xdist
 wheel
-numpy<1.23.0
+numpy
path: build/test-requirements.txt
text: @user1:
Remove this? setup.py should cover it

@author:
I don't understand - we want the tests to run on the most recent version of numpy.

@user1:
Right. There is no upper bound on version constraint on numpy.

This numpy was not there before I added the `<1.23.0` and we still hit that error. So I think removing this should be fine.

@author:
Oh I see, you meant remove the line, not remove the change. Will do.