Skip to content
This repository was archived by the owner on Oct 31, 2025. It is now read-only.

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

Merged
DanPuzzuoli merged 5 commits into from
Dec 13, 2022
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
3dee1b7
adding jax.Array to the JAX types in dispatch
DanPuzzuoli Dec 12, 2022
1bb11ac
forgot to commit the actual required change
DanPuzzuoli Dec 13, 2022
2f531b6
adding comment making it clear why we're checkign hasattr(jax, 'Array')
DanPuzzuoli Dec 13, 2022
77d13f9
making the check more in lien with previous existing code
DanPuzzuoli Dec 13, 2022
f482420
fixing JAX version number
DanPuzzuoli Dec 13, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions qiskit_dynamics/dispatch/backends/jax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions releasenotes/notes/jax-4-compatibility-8e3398e95f758dfe.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
---
upgrade:
- |
The ``jax.Array`` class has been added to the dispatcher for compatibility with JAX 0.4.0.