Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add warning for JAX versions on import of Dynamics #232

Merged
merged 10 commits into from
Jun 17, 2023
Merged

Add warning for JAX versions on import of Dynamics #232

merged 10 commits into from
Jun 17, 2023

Conversation

DanPuzzuoli
Copy link
Collaborator

@DanPuzzuoli DanPuzzuoli commented Jun 16, 2023 •
edited
Loading

Summary

Closes #231

This PR updates the JAX type registration in dispatch, and also adds a warning at JAX import time about the version bounds for JAX.

Details and comments

The docs build has also been set to run with JAX version 0.4.3 to avoid this warning.

@DanPuzzuoli DanPuzzuoli mentioned this pull request Jun 16, 2023
Closed
@DanPuzzuoli DanPuzzuoli requested a review from wshanks June 16, 2023 19:45
wshanks
wshanks previously approved these changes Jun 16, 2023
View reviewed changes
Copy link
Contributor

@wshanks wshanks left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had a couple suggestions but this looks fine to me (other than the tracer type changes which I am not sure about).

qiskit_dynamics/dispatch/backends/jax.py Outdated Show resolved Hide resolved
qiskit_dynamics/dispatch/backends/jax.py Outdated Show resolved Hide resolved
tox.ini Outdated Show resolved Hide resolved
@DanPuzzuoli
Copy link
Collaborator Author

DanPuzzuoli commented Jun 16, 2023 •
edited
Loading

@wshanks I've implemented your suggestion for the warning to not trigger if the JAX version is 0.4.4, 0.4.5, or 0.4.6 and the user has already set the os environment variable. I've reverted the docs to use JAX 0.4.6 as the warning will no longer be triggered.

@DanPuzzuoli
Copy link
Collaborator Author

So I still ended up needing to set os.environ["JAX_JIT_PJIT_API_MERGE"] = "0" in docs/conf.py to avoid the warning in the docs build, but luckily that worked. I've put that in with a comment, and have added a description of this PR to issue #190 (where I'm keeping track of all the little changes that we've had to make because of these JAX issues).

I was actually also able to completely remove adding DeviceArray to JAX_TYPES - Array and Tracer turn out to be all that's necessary these days.

@DanPuzzuoli DanPuzzuoli mentioned this pull request Jun 17, 2023
Merged
wshanks
wshanks approved these changes Jun 17, 2023
View reviewed changes
Copy link
Contributor

@wshanks wshanks left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good!

@DanPuzzuoli DanPuzzuoli merged commit 37d12ac into qiskit-community:main Jun 17, 2023
DanPuzzuoli added a commit to DanPuzzuoli/qiskit-dynamics that referenced this pull request Jun 20, 2023
@DanPuzzuoli
Add warning for JAX versions on import of Dynamics (qiskit-community#232
360ab09 
)
@DanPuzzuoli DanPuzzuoli mentioned this pull request Jun 20, 2023
Merged
DanPuzzuoli added a commit to DanPuzzuoli/qiskit-dynamics that referenced this pull request Aug 2, 2023
@DanPuzzuoli
Add warning for JAX versions on import of Dynamics (qiskit-community#232
6ca6293 
)
DanPuzzuoli added a commit that referenced this pull request Aug 8, 2023
@DanPuzzuoli @to24toro
0.4.2 patch release (#234)
d217616 
* adding sympy as an explicit requirement (#233)

* Add warning for JAX versions on import of Dynamics (#232)

* Add probability normalization to DynamicsBackend sampling routine (#239)

* Add warning if digital carrier exceeds Nyquist frequency in pulse -> signal conversion (#242)

* Fix bug with carrier_freq being a JAX tracer if envelope is constant in Signal (#247)

* Remove subsystem_labels option from DynamicsBackend (#248)

* Rename subsystem_dims Dict to subsystem_dims_dict (#250)

* Fix ClassicalRegister counting in DynamicsBackend (#252)

* temporary fix for docs (#253)

* Drop support for backendV2 in from_backend (#249)

* incrementing version number

* cleaning up reno files

---------

Co-authored-by: Kento Ueda <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@wshanks wshanks wshanks approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

Incompatibility with jax 0.4.11
2 participants
@DanPuzzuoli @wshanks