qml.wire.Wires accepts JAX arrays
#6312
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Hello. You may have forgotten to update the changelog!
|
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## master #6312 +/- ##
=======================================
Coverage 99.39% 99.39%
=======================================
Files 445 445
Lines 42286 42315 +29
=======================================
+ Hits 42031 42060 +29
Misses 255 255 ☔ View full report in Codecov by Sentry. |
albi3ro
reviewed
Oct 2, 2024
dime10
reviewed
Oct 2, 2024
albi3ro
reviewed
Oct 2, 2024
…o Wires_accept_JAX_array
dime10
approved these changes
Oct 3, 2024
…o Wires_accept_JAX_array
austingmhuang
pushed a commit
that referenced
this pull request
Oct 23, 2024
**Context:** In JAX JIT/plxpr/qjit workflows, it is quite easy for wires to end up with jax array values. JAX arrays are not hashable and are rejected. Still, we can treat them as iterable and ensure they are accepted as wires. **Description of the Change:** We use the `to_list` and `item` methods of `JAX` arrays as a trick to make them tuples before hashing them. **Benefits:** `qml.wire.Wires` can accept `JAX` arrays as input. **Possible Drawbacks:** The change should only be relevant when `jax.numpy.array`s are explicitly provided. Therefore, in principle, we should not face problems with the rest of the code. **Related GitHub Issues:** None. **Related Shortcut Stories** [sc-72593] [sc-74904]
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Context: In JAX JIT/plxpr/qjit workflows, it is quite easy for wires to end up with jax array values. JAX arrays are not hashable and are rejected. Still, we can treat them as iterable and ensure they are accepted as wires.
Furthermore, if we try to use JAX tracers as wires in PennyLane using JAX 0.4.30+, we currently get FutureWarning in some of the tests:
This future warning was introduced in JAX 0.4.30 and is still not implemented as an error in 0.4.32, but both PL and Catalyst currently use version 0.4.28.
Description of the Change: We use the
to_listanditemmethods ofJAXarrays as a trick to make them tuples before hashing them to ensure that they are accepted as iterable by theWireclass. Furthermore, we decided to implement in PL the same patch adopted in Catalyst to avoid the warning reported above as temporary solution.Benefits:
qml.wire.Wirescan acceptJAXarrays as input.Possible Drawbacks: The first change should only be relevant when
jax.numpy.arrays are explicitly provided. Therefore, we should not face problems with the rest of the code.As for the second change, it is highly recommended to re-visit the way JAX tracers are handled in the PL pipeline since in a future JAX version we will not be able to hash tracers anymore.
Related GitHub Issues: None.
Related Shortcut Stories [sc-72593] [sc-74904]