Current JAX release incompatible with numpy 1.21

[jakevdp]

Numpy 1.21 was released a few hours ago, and appears incompatible with the current jaxlib release.

[shoyer]

numpy/numpy#19305 is the source of one of the remaining errors.

[hawkinsp]

I've done all I plan to do on this for now: remaining is to release a new jaxlib (in progress) and to do something to fix or work around the numpy bug.

[jakevdp]

Alright, I'll self-assign so I can close the issue after the next jaxlib release.

[hawkinsp]

Note that we can't close it until we do something to work around the numpy problem. It could just be skipping those tests under NumPy 1.21 for now.

If you want to try running the tests before the new jaxlib lands, apply the following patch:

diff --git a/jax/__init__.py b/jax/__init__.py
index 17629466..caca87a4 100644
--- a/jax/__init__.py
+++ b/jax/__init__.py
@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

+import numpy as np
+np.typeDict = np.sctypeDict
+
 # Set default logging level before any logging happens.
 import os as _os
 _os.environ.setdefault('TF_CPP_MIN_LOG_LEVEL', '1')

[jakevdp]

Part of the action item here may be to add a CI job to detect these kinds of breakages before release. @shoyer pointed out the approach that xarray uses, which is a github action that runs tests against upstream dev versions and opens a github issue in case of failure: https://github.com/pydata/xarray/blob/main/.github/workflows/upstream-dev-ci.yaml