Extend jax.vmap to the multidimensional case
#6422
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
…o vmap_program_capture
…com/PennyLaneAI/pennylane into vmap_program_capture
…o vmap_program_capture
…o vmap_program_capture
…o vmap_program_capture
…o vmap_program_capture
…/pennylane into experim_param_broad_capture
…o experim_param_broad_capture
…o vmap_program_capture
…o experim_param_broad_capture
…yLaneAI/pennylane into vmap_program_capture
…o vmap_program_capture
…/pennylane into vmap_multidimensional_input
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## master #6422 +/- ##
==========================================
- Coverage 99.45% 99.45% -0.01%
==========================================
Files 450 450
Lines 42088 42085 -3
==========================================
- Hits 41857 41854 -3
Misses 231 231 ☔ View full report in Codecov by Sentry. 🚨 Try these New Features:
|
…o vmap_program_capture
…/pennylane into vmap_multidimensional_input
…/pennylane into vmap_multidimensional_input
albi3ro
added a commit
that referenced
this pull request
Nov 8, 2024
**Context:** This PR is the first step to enable the usage of `jax.vmap`
with quantum circuits when `qml.capture.enabled()` is `True`.
**Description of the Change:** We implemented a batching rule for the
captured `qnode` and modified the abstract evaluation to keep track of
the batch dimension.
**Benefits:** Allows vectorization (although with several limitations at
this stage) with `qml.capture.enabled()` via `jax.vmap`
**Possible Drawbacks:** Right now, there are 2 main limitations that I
see.
- Limitation 1: *Multidimensional arrays*
Right now, the implementation does not work with multidimensional
arrays. For example, the following does not work:
```
qml.capture.enable()
@qml.qnode(qml.device("default.qubit", wires=2))
def circuit(x):
qml.RX(x[1], 0)
return qml.expval(qml.Z(0))
# jax.vmap(circuit)(jax.numpy.array([[0.1, 0.2], [0.3, 0.4]])) Generates an error
```
This limitation will be removed in the [following
PR](#6422).
- Limitation 2 *How to prevent the user from bypassing `jax.vmap` ?*
```
qml.capture.enable()
@qml.qnode(qml.device("default.qubit", wires=2))
def circuit(x, y):
qml.RX(x, 0)
qml.RY(y, 0)
return qml.expval(qml.Z(0))
x = jax.numpy.array([0.1, 0.2])
y = jax.numpy.array([0.1, 0.2])
jax.vmap(circuit, in_axes=(0, None))(x, y)
```
This works by accident, but the second argument is vectorized along 2
dimensions, although it shouldn't!
With `in_axes=(0, None)` we should raise an error if `qml.RY` receives
something that is not a scalar (as [it happens in
Catalyst](https://github.com/PennyLaneAI/catalyst/blob/7c5b828d5173cdaa52073d30a5f3a7df660b37d6/frontend/catalyst/jax_primitives.py#L1206)).
At this stage, we simply raise a warning because we don't have a way to
check this inside the QNode batching rule. To fix this behavior, I think
we need to add more properties to the `AbstractOperator` class and
improve the integration with the captured `QNode`. This is most probably
also necessary for capturing parameter broadcasting in PL.
**Related GitHub Issues:** None.
**Related Shortcut Stories:** [sc-73779] [sc-73782] [sc-73783]
---------
Co-authored-by: Christina Lee <[email protected]>
…o vmap_multidimensional_input
mudit2812
pushed a commit
that referenced
this pull request
Nov 11, 2024
**Context:** This PR is the first step to enable the usage of `jax.vmap`
with quantum circuits when `qml.capture.enabled()` is `True`.
**Description of the Change:** We implemented a batching rule for the
captured `qnode` and modified the abstract evaluation to keep track of
the batch dimension.
**Benefits:** Allows vectorization (although with several limitations at
this stage) with `qml.capture.enabled()` via `jax.vmap`
**Possible Drawbacks:** Right now, there are 2 main limitations that I
see.
- Limitation 1: *Multidimensional arrays*
Right now, the implementation does not work with multidimensional
arrays. For example, the following does not work:
```
qml.capture.enable()
@qml.qnode(qml.device("default.qubit", wires=2))
def circuit(x):
qml.RX(x[1], 0)
return qml.expval(qml.Z(0))
# jax.vmap(circuit)(jax.numpy.array([[0.1, 0.2], [0.3, 0.4]])) Generates an error
```
This limitation will be removed in the [following
PR](#6422).
- Limitation 2 *How to prevent the user from bypassing `jax.vmap` ?*
```
qml.capture.enable()
@qml.qnode(qml.device("default.qubit", wires=2))
def circuit(x, y):
qml.RX(x, 0)
qml.RY(y, 0)
return qml.expval(qml.Z(0))
x = jax.numpy.array([0.1, 0.2])
y = jax.numpy.array([0.1, 0.2])
jax.vmap(circuit, in_axes=(0, None))(x, y)
```
This works by accident, but the second argument is vectorized along 2
dimensions, although it shouldn't!
With `in_axes=(0, None)` we should raise an error if `qml.RY` receives
something that is not a scalar (as [it happens in
Catalyst](https://github.com/PennyLaneAI/catalyst/blob/7c5b828d5173cdaa52073d30a5f3a7df660b37d6/frontend/catalyst/jax_primitives.py#L1206)).
At this stage, we simply raise a warning because we don't have a way to
check this inside the QNode batching rule. To fix this behavior, I think
we need to add more properties to the `AbstractOperator` class and
improve the integration with the captured `QNode`. This is most probably
also necessary for capturing parameter broadcasting in PL.
**Related GitHub Issues:** None.
**Related Shortcut Stories:** [sc-73779] [sc-73782] [sc-73783]
---------
Co-authored-by: Christina Lee <[email protected]>
rmoyard
approved these changes
Nov 18, 2024
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: This PR extends the captured
jax.vmapversion to the multidimensional input case. For further details, we refer to the description of the first PR.Description of the Change: As above. For 'multidimensional input case' we mean something like the following:
Benefits: Now
jax.vmapcan be used with captured enabled if the input parameter is an array with a shape greater than 1.Possible Drawbacks: None that I can think of.
Related GitHub Issues: None.
Related Shortcut Stories: [sc-76381]