Add fixes for jax 0.4 (#601)

* Add fixes for jax 0.4

* Add restriction

* Require ci to test jax with python 3.7

.github/workflows/ci.yml

@@ -67,7 +67,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7]
+        python-version: [3.7,3.8]
     env:
       CI: 1
       FUNSOR_BACKEND: jax

docs/requirements.txt

@@ -1,3 +1,4 @@
+ipython<=8.6.0  # restrict for https://github.com/ipython/ipython/issues/13845
 makefun
 multipledispatch
 nbsphinx==0.8.1

funsor/jax/__init__.py

@@ -1,8 +1,10 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+import re
+
 from jax.core import Tracer
-from jax.interpreters.xla import DeviceArray
+from jax.numpy import ndarray
 
 from funsor.tensor import tensor_to_funsor
 from funsor.terms import to_funsor
@@ -14,15 +16,23 @@
 del _  # flake8
 
 
-to_funsor.register(DeviceArray)(tensor_to_funsor)
+to_funsor.register(ndarray)(tensor_to_funsor)
 to_funsor.register(Tracer)(tensor_to_funsor)
 
 
-@quote.register(DeviceArray)
+@quote.register(ndarray)
 def _quote(x, indent, out):
     """
-    Work around JAX's DeviceArray not supporting reproducible repr.
+    Work around JAX's ndarray not supporting reproducible repr.
     """
+    # After JAX 0.4, jnp.ones(3) is no longer a DeviceArray, but an ndarray.
+    # In addition, a tracer is also an ndarray - so we need to handler it
+    # separately here.
+    if isinstance(x, Tracer):
+        # Default implementation.
+        line = re.sub("\n\\s*", " ", repr(x))
+        out.append((indent, line))
+        return
     if x.size >= quote.printoptions["threshold"]:
         data = "..." + " x ".join(str(d) for d in x.shape) + "..."
     else: