Adding jax.Array to the JAX types in dispatch (#162)

qiskit-community · Dec 13, 2022 · 9c64ad1 · 9c64ad1
1 parent 44a04bd
commit 9c64ad1
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/qiskit_dynamics/dispatch/backends/jax.py b/qiskit_dynamics/dispatch/backends/jax.py
@@ -24,6 +24,14 @@
 
     JAX_TYPES = (DeviceArray, Tracer, JaxprTracer, JVPTracer)
 
+    try:
+        # This class was introduced in 0.4.0
+        from jax import Array
+
+        JAX_TYPES += (Array,)
+    except ImportError:
+        pass
+
     try:
         # This class is not in older versions of Jax
         from jax.interpreters.partial_eval import DynamicJaxprTracer

diff --git a/releasenotes/notes/jax-4-compatibility-8e3398e95f758dfe.yaml b/releasenotes/notes/jax-4-compatibility-8e3398e95f758dfe.yaml
@@ -0,0 +1,4 @@
+---
+upgrade:
+  - |
+    The ``jax.Array`` class has been added to the dispatcher for compatibility with JAX 0.4.0.