[Program Capture] Capture & execute `qml.grad` in plxpr #6120

dwierichs · 2024-08-21T12:37:55Z

Context:
The new qml.capture module does not support differentiation yet.

Description of the Change:
This PR takes the first step towards differentiability in plxpr.
It adds the capability of capturing qml.grad as a "nested jaxpr" primitive.
When executing the captured program, qml.grad is essentially changed to jax.grad, because executing Autograd autodifferentiation within the Jaxpr ecosystem is not sensible.

Benefits:
Capture first differentiation instructions

Possible Drawbacks:
The current implementation requires a jvp construction for every evaluation of a QNode gradient. This means that this JVP function is reconstructed for every evaluation call, if I'm not mistaken, making the code significantly less performant with capture than without. Of course, the longer term plan is to process the plxpr into lower-level code by lowering the grad primitive itself, in which case this problem goes away.
A similar redundancy is implemented in QNode: Whenever a qnode primitive is evaluated, a new QNode is created (and only ever evaluated once). This disables caching, for example, unless a cache is passed around explicitly.

Related GitHub Issues:

[sc-71858]

github-actions · 2024-08-21T12:38:10Z

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

codecov · 2024-08-21T13:12:15Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.60%. Comparing base (df63953) to head (db99812).
Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6120   +/-   ##
=======================================
  Coverage   99.60%   99.60%           
=======================================
  Files         445      446    +1     
  Lines       42364    42423   +59     
=======================================
+ Hits        42198    42257   +59     
  Misses        166      166

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

dwierichs

Just leaving some (hopefully helpful) comments.

pennylane/_grad.py

pennylane/capture/capture_qnode.py

pennylane/capture/primitives.py

mudit2812

🎉

[Program Capture] Capture & execute `qml.grad` in plxpr

**Context:** We're adding support for differentiation in plxpr, also see #6120. **Description of the Change:** This PR adds support for `qml.jacobian`, similar to the support for `qml.grad`. Note that Pytree support will be needed to allow for multi-argument derivatives. **Benefits:** Capture derivatives of non-scalar functions. **Possible Drawbacks:** See discussion around `qml.grad` in #6120. **Related GitHub Issues:** [sc-71860] --------- Co-authored-by: Christina Lee <christina@xanadu.ai>

**Context:** The new `qml.capture` module does not support differentiation yet. **Description of the Change:** This PR takes the first step towards differentiability in plxpr. It adds the capability of capturing `qml.grad` as a "nested jaxpr" primitive. When executing the captured program, `qml.grad` is essentially changed to `jax.grad`, because executing Autograd autodifferentiation within the Jaxpr ecosystem is not sensible. **Benefits:** Capture first differentiation instructions **Possible Drawbacks:** The current implementation requires a `jvp` construction for every evaluation of a QNode gradient. This means that this JVP function is reconstructed for every evaluation call, if I'm not mistaken, making the code significantly less performant with `capture` than without. Of course, the longer term plan is to process the plxpr into lower-level code by lowering the `grad` primitive itself, in which case this problem goes away. A similar redundancy is implemented in `QNode`: Whenever a `qnode` primitive is evaluated, a new `QNode` is created (and only ever evaluated once). This disables caching, for example, unless a cache is passed around explicitly. **Related GitHub Issues:** [sc-71858] --------- Co-authored-by: Christina Lee <christina@xanadu.ai>

**Context:** We're adding support for differentiation in plxpr, also see #6120. **Description of the Change:** This PR adds support for `qml.jacobian`, similar to the support for `qml.grad`. Note that Pytree support will be needed to allow for multi-argument derivatives. **Benefits:** Capture derivatives of non-scalar functions. **Possible Drawbacks:** See discussion around `qml.grad` in #6120. **Related GitHub Issues:** [sc-71860] --------- Co-authored-by: Christina Lee <christina@xanadu.ai>

**Context:** The new `qml.capture` module does not support differentiation yet. **Description of the Change:** This PR takes the first step towards differentiability in plxpr. It adds the capability of capturing `qml.grad` as a "nested jaxpr" primitive. When executing the captured program, `qml.grad` is essentially changed to `jax.grad`, because executing Autograd autodifferentiation within the Jaxpr ecosystem is not sensible. **Benefits:** Capture first differentiation instructions **Possible Drawbacks:** The current implementation requires a `jvp` construction for every evaluation of a QNode gradient. This means that this JVP function is reconstructed for every evaluation call, if I'm not mistaken, making the code significantly less performant with `capture` than without. Of course, the longer term plan is to process the plxpr into lower-level code by lowering the `grad` primitive itself, in which case this problem goes away. A similar redundancy is implemented in `QNode`: Whenever a `qnode` primitive is evaluated, a new `QNode` is created (and only ever evaluated once). This disables caching, for example, unless a cache is passed around explicitly. **Related GitHub Issues:** [sc-71858] --------- Co-authored-by: Christina Lee <christina@xanadu.ai>

**Context:** We're adding support for differentiation in plxpr, also see #6120. **Description of the Change:** This PR adds support for `qml.jacobian`, similar to the support for `qml.grad`. Note that Pytree support will be needed to allow for multi-argument derivatives. **Benefits:** Capture derivatives of non-scalar functions. **Possible Drawbacks:** See discussion around `qml.grad` in #6120. **Related GitHub Issues:** [sc-71860] --------- Co-authored-by: Christina Lee <christina@xanadu.ai>

…jacobian` (#6134) **Context:** #6120 and #6127 add support to capture `qml.grad` and `qml.jacobian` in plxpr. Once captured, they dispatch to `jax.grad` and `jax.jacobian`. **Description of the Change:** This PR adds support for pytree inputs and outputs of the differentiated functions, similar to #6081. For this, it extends the internal class `FlatFn` by the extra functionality to turn the wrapper into a `*flat_args -> *flat_outputs` function, instead of a `*pytree_args -> *flat_outputs` function. **Benefits:** Pytree support 🌳 **Possible Drawbacks:** **Related GitHub Issues:** [sc-70930] [sc-71862] --------- Co-authored-by: Christina Lee <christina@xanadu.ai> Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>

**Context:** We're adding support for differentiation in plxpr, also see #6120. **Description of the Change:** This PR adds support for `qml.jacobian`, similar to the support for `qml.grad`. Note that Pytree support will be needed to allow for multi-argument derivatives. **Benefits:** Capture derivatives of non-scalar functions. **Possible Drawbacks:** See discussion around `qml.grad` in #6120. **Related GitHub Issues:** [sc-71860] --------- Co-authored-by: Christina Lee <christina@xanadu.ai>

…jacobian` (#6134) **Context:** #6120 and #6127 add support to capture `qml.grad` and `qml.jacobian` in plxpr. Once captured, they dispatch to `jax.grad` and `jax.jacobian`. **Description of the Change:** This PR adds support for pytree inputs and outputs of the differentiated functions, similar to #6081. For this, it extends the internal class `FlatFn` by the extra functionality to turn the wrapper into a `*flat_args -> *flat_outputs` function, instead of a `*pytree_args -> *flat_outputs` function. **Benefits:** Pytree support 🌳 **Possible Drawbacks:** **Related GitHub Issues:** [sc-70930] [sc-71862] --------- Co-authored-by: Christina Lee <christina@xanadu.ai> Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>

dwierichs added 4 commits August 21, 2024 13:22

first prototype

206ca6a

first tests

a466a83

cleanup

8b402f7

-a

9e8f4c9

changelog, lint

cfae262

add tests...

ed8f515

dwierichs changed the title ~~Capture qml.grad in plxpr~~ [Program Capture] Capture qml.grad in plxpr Aug 21, 2024

dwierichs changed the title ~~[Program Capture] Capture qml.grad in plxpr~~ [Program Capture] Capture & execute qml.grad in plxpr Aug 22, 2024

dwierichs added 6 commits August 22, 2024 10:01

move primitive

79fe790

import

b90fa2b

lint

b9b6536

lint

ff54b86

parshift test

fde341f

prepare jac

3cd33f1

dwierichs mentioned this pull request Aug 22, 2024

[Program Capture] Capture & execute qml.jacobian in plxpr #6127

Merged

dwierichs added 5 commits August 22, 2024 13:27

make grad differentiable

a087e53

nested grad test

393aa35

lint

f1a6b4d

Merge branch 'master' into capture-grad

b86619b

lint

f678a91

dwierichs added the review-ready 👌 PRs which are ready for review by someone from the core team. label Aug 22, 2024

mudit2812 self-requested a review August 22, 2024 14:22

albi3ro self-requested a review August 22, 2024 14:37

dwierichs added 2 commits August 22, 2024 17:18

Merge branch 'master' into capture-grad

15ed932

Merge branch 'master' into capture-grad

692db13

dwierichs mentioned this pull request Aug 23, 2024

[Program Capture] Add pytree support to captured qml.grad and qml.jacobian #6134

Merged

dwierichs commented Aug 26, 2024

View reviewed changes

pennylane/_grad.py Outdated Show resolved Hide resolved

pennylane/_grad.py Show resolved Hide resolved

pennylane/capture/capture_qnode.py Show resolved Hide resolved

pennylane/capture/primitives.py Outdated Show resolved Hide resolved

Merge branch 'master' into capture-grad

0e407c3

dwierichs added 2 commits September 5, 2024 18:08

while_loop

552417b

import fix

81ab600

dwierichs requested a review from mudit2812 September 5, 2024 16:11

mudit2812 approved these changes Sep 5, 2024

View reviewed changes

dwierichs added 6 commits September 5, 2024 19:33

lint

1af5b97

fix import

c3fbd78

import and skip order

e73f2fb

Merge branch 'master' into capture-grad

3d2cb89

[skip ci]

724675f

Merge branch 'master' into capture-grad

9fda181

dwierichs enabled auto-merge September 9, 2024 08:46

lint

db99812

dwierichs added this pull request to the merge queue Sep 9, 2024

github-merge-queue bot pushed a commit that referenced this pull request Sep 9, 2024

Merge pull request #6120 from PennyLaneAI/capture-grad

e63d7f8

[Program Capture] Capture & execute `qml.grad` in plxpr

github-merge-queue bot removed this pull request from the merge queue due to no response for status checks Sep 9, 2024

dwierichs added this pull request to the merge queue Sep 9, 2024

github-merge-queue bot removed this pull request from the merge queue due to no response for status checks Sep 9, 2024

dwierichs merged commit 1d220cf into master Sep 9, 2024
37 checks passed

dwierichs deleted the capture-grad branch September 9, 2024 13:06

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Program Capture] Capture & execute `qml.grad` in plxpr #6120

[Program Capture] Capture & execute `qml.grad` in plxpr #6120

dwierichs commented Aug 21, 2024 •

edited

Loading

github-actions bot commented Aug 21, 2024

codecov bot commented Aug 21, 2024 •

edited

Loading

dwierichs left a comment

mudit2812 left a comment

[Program Capture] Capture & execute qml.grad in plxpr #6120

[Program Capture] Capture & execute qml.grad in plxpr #6120

Conversation

dwierichs commented Aug 21, 2024 • edited Loading

github-actions bot commented Aug 21, 2024

codecov bot commented Aug 21, 2024 • edited Loading

Codecov Report

dwierichs left a comment

Choose a reason for hiding this comment

mudit2812 left a comment

Choose a reason for hiding this comment

[Program Capture] Capture & execute `qml.grad` in plxpr #6120

[Program Capture] Capture & execute `qml.grad` in plxpr #6120

dwierichs commented Aug 21, 2024 •

edited

Loading

codecov bot commented Aug 21, 2024 •

edited

Loading