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

[Capture] Add backprop validation #6852

Merged
merged 9 commits into from
Jan 21, 2025
Merged

[Capture] Add backprop validation #6852

merged 9 commits into from
Jan 21, 2025

Conversation

albi3ro
Copy link
Contributor

@albi3ro albi3ro commented Jan 17, 2025

Context:

We currently use un-validated backprop for differentiation with program capture. This leads to some unintuitive errors if you try and take a gradient on lightning with capture enabled.

Description of the Change:

Adds some validation to make sure the device supports backprop. Adds the backprop logic to a _backprop jvp function, and dispatches to that method based on the diff method.

Benefits:

Improved error messages when backprop or the requested diff method isn't supported.

Possible Drawbacks:

The code currently is a little clunky, but it is private so we should be able to move things around once we have more information.

Related GitHub Issues:

[sc-82166]

@github-actions GitHub Actions
Copy link
Contributor

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 Codecov
Copy link

codecov bot commented Jan 17, 2025 •
edited
Loading

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.59%. Comparing base (98bb29b) to head (fc6b28f).
Report is 1 commits behind head on master.

Additional details and impacted files 
@@            Coverage Diff             @@
##           master    #6852      +/-   ##
==========================================
- Coverage   99.60%   99.59%   -0.01%     
==========================================
  Files         476      476              
  Lines       45182    45193      +11     
==========================================
+ Hits        45002    45012      +10     
- Misses        180      181       +1

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

Base automatically changed from capture-qnode-reorg to master January 17, 2025 21:00
andrijapau
andrijapau approved these changes Jan 21, 2025
View reviewed changes
Copy link
Contributor

@andrijapau andrijapau 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 to me! 🙂

@mudit2812 mudit2812 self-requested a review January 21, 2025 19:37
mudit2812
mudit2812 approved these changes Jan 21, 2025
View reviewed changes
Copy link
Contributor

@mudit2812 mudit2812 left a comment

Choose a reason for hiding this comment

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

Thanks Christina! I just have one non-blocking comment :)

pennylane/workflow/_capture_qnode.py Show resolved Hide resolved
@albi3ro albi3ro enabled auto-merge (squash) January 21, 2025 21:05
@albi3ro albi3ro merged commit 90dc57c into master Jan 21, 2025
46 checks passed
@albi3ro albi3ro deleted the setup-jvp-structure branch January 21, 2025 21:21
mudit2812 added a commit that referenced this pull request Jan 24, 2025
@mudit2812
Squashed commit of the following:
25c12e8 
commit 666941f
Author: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com>
Date:   Fri Jan 24 11:22:28 2025 -0500

    `qml.execute` no longer accepts `mcm_config` argument (#6807)

    **Context:**

    Further `qml.workflow` clean-up. This enables `qml.execute` to mimic the
    signature of `QNode`. This ensures that we don't get incompatible
    configurations because we have two different entry points.

    **Description of the Change:**
    - [x] Catalyst: PennyLaneAI/catalyst#1452
    - [x] Lightning: No instances of deprecated code found.
    - [x] QML Demos: No instances of deprecated code found.
    - [x] Pennylane-AQT: No instances of deprecated code found.
    - [x] Pennylane-Qiskit: No instances of deprecated code found.
    - [x] Pennylane-IonQ: No instances of deprecated code found.
    - [x] Pennylane-Qrack: No instances of deprecated code found.
    - [x] Pennylane-Cirq: No instances of deprecated code found.
    - [x] Pennylane-Qulacs: No instances of deprecated code found.

    Introduce kwargs `postselect_mode` and `mcm_method` to `qml.execute`
    signature. Raise deprecation warning to user if they try to use
    `mcm_config` kwarg.

    Side-effect: Had to `xfail` any tests that had to do with Catalyst since
    they assume certain keys from the `QNode.execute_kwargs`. Will be
    reverted in #6873.

    **Benefits:**

    Keyword parity with `qml.QNode`.

    **Possible Drawbacks:** None identified.

    [sc-80541]

    ---------

    Co-authored-by: Christina Lee <christina@xanadu.ai>
    Co-authored-by: Yushao Chen (Jerry) <chenys13@outlook.com>

commit dbb3315
Author: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai>
Date:   Fri Jan 24 09:51:39 2025 +0000

    [no ci] bump nightly version

commit f367e01
Author: Rashid N H M <95639609+rashidnhm@users.noreply.github.com>
Date:   Thu Jan 23 10:47:17 2025 -0500

    Add merge queue trigger for required workflows (#6860)

    **Context:**
    This pull request adds a new trigger to the existing workflows that run
    `on.pull_request`. This trigger indicates to GitHub which workflows
    needs to be run when a merge queue is building.

    **Description of the Change:**
    The change is adding `on.merge_group` to our required workflows.

    **Benefits:**
    This change itself will not enable merge queues, that needs to be
    enabled from the admin settings of branch protection rules.

    The changes in this PR mainly tell merge queues which workflows to run.

    **Possible Drawbacks:**
    Usage of merge queue is a new thing for pennylane, if issues arise we
    can rollback.

    **Related GitHub Issues:**
    None. [sc-82039]

commit 8a12fa5
Author: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com>
Date:   Thu Jan 23 10:06:38 2025 -0500

    Promote `gradient_kwargs` to a positional keyword argument in `QNode` (#6828)

    **Context:**
    `gradient_kwargs` is now a positional keyword argument for the `QNode`.
    This means you can not simply express,
    ```python
    qml.QNode(func, dev, h=1)
    ```
    instead, you must deliberately,
    ```python
    qml.QNode(func, dev, gradient_kwargs={"h":1})
    ```
    This allows easier and cleaner input validation.

    This PR could have wide-spread impact as it is very common to just
    specify `gradient_kwargs` casually as additional kwargs.
    - [x] Catalyst: PennyLaneAI/catalyst#1480
    - [x] Lightning:
    PennyLaneAI/pennylane-lightning#1045
    - [x] QML Demos: No instances of deprecated code found.
    - [x] Pennylane-AQT: No instances of deprecated code found.
    - [x] Pennylane-Qiskit: No instances of deprecated code found.
    - [x] Pennylane-IonQ: No instances of deprecated code found.
    - [x] Pennylane-Qrack: No instances of deprecated code found.
    - [x] Pennylane-Cirq: No instances of deprecated code found.
    - [x] Pennylane-Qulacs: No instances of deprecated code found.

    **Description of the Change:**

    Allow additional kwargs for now to ensure same functionality, but raise
    a deprecation warning. Append those additional kwargs to the internal
    gradient_kwargs dictionary.

    **Benefits:** Improved input validation for users.

    **Possible Drawbacks:** Might have missed some eco-system changes.
    Especially with CI **sometimes** not raising
    `PennyLaneDeprecationWarning`s as errors 😢 .

    [sc-81531]

    ---------

    Co-authored-by: Christina Lee <christina@xanadu.ai>

commit 875ae11
Author: Christina Lee <christina@xanadu.ai>
Date:   Thu Jan 23 09:31:47 2025 -0500

    Revert end-to-end jitting with default qubit (#6869)

    **Context:**

    In #6788 , we started allowing executions on default qubit to be jitted
    from end-to-end. Unfortunately, we found that the compilation overheads
    on these executions can get very, very expensive. So until we find a way
    to reduce the compilation overheads, we are using pure callbacks and
    conversion to numpy.

    **Description of the Change:**

    Default to `convert_to_numpy=False`, and xfail relevant tests. This
    change can be undone once we figure out how to resolve the compilation
    issue.

    **Benefits:**

    Reduced compilation overheads, because the execution itself does not get
    compiled.

    **Possible Drawbacks:**

    Slow down on post-compiled workflows. No way to jit an entire execution
    on default qubit.

    **Related GitHub Issues:**

    ---------

    Co-authored-by: Pietropaolo Frisoni <pietropaolo.frisoni@xanadu.ai>

commit 63cca88
Author: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai>
Date:   Thu Jan 23 09:51:30 2025 +0000

    [no ci] bump nightly version

commit 61dbc71
Author: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com>
Date:   Wed Jan 22 16:15:54 2025 -0500

    Deprecate `qml.gradients.hamiltonian_grad` (#6849)

    **Context:**

    **Description of the Change:**

    _Source-Code_

    Standard deprecation of `hamiltonian_grad` function.

    _Test suite_

    I noticed that the tests I've removed from `test_parameter_shift.py`
    have improved duplicates in `tests/workflow/interfaces` under the
    `TestHamiltonianWorkflows` test class. Therefore, they were all removed
    except `test_jax`. The reason is that this test follows the outdated
    workflow that still hits the branch in
    `parameter_shift.py::expval_param_shift` that raises the deprecation
    warning. So, I've added a warning and left that test. This should be
    removed with the `hamiltonian_grad` function next release.

    **Impact:**

    No deprecated code found elsewhere. Impact to the eco-system should be
    minimal.

    [sc-81526]

commit 3e1521b
Author: Mudit Pandey <mudit.pandey@xanadu.ai>
Date:   Wed Jan 22 11:09:07 2025 -0500

    [Capture] Add a `QmlPrimitive` class to differentiate between different types of primitives (#6847)

    This PR adds a `QmlPrimitive` subclass of `jax.core.Primitive`. This
    class contains a `prim_type` property set using a new `PrimitiveType`
    enum. `PrimitiveType`s currently available are "default", "operator",
    "measurement", "transform", and "higher_order". This can be made more or
    less fine grained as needed, but should be enough to differentiate
    between different types of primitives for now. Additionally, this PR:

    * updates `NonInterpPrimitive` to be a subclass of `QmlPrimitive`
    * updates all existing PennyLane primitives to be either `QmlPrimitive`
    or `NonInterpPrimitive`. See [this
    comment](#6851 (comment))
    to see the logic used to determine which `Primitive` subclass is used
    for each primitive.
    * updates `PlxprInterpreter.eval` and `CancelInversesInterpreter.eval`
    to use this `prim_type` property.

    [sc-82420]

    ---------

    Co-authored-by: Pietropaolo Frisoni <pietropaolo.frisoni@xanadu.ai>

commit fdf34ec
Author: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai>
Date:   Wed Jan 22 09:51:47 2025 +0000

    [no ci] bump nightly version

commit 90dc57c
Author: Christina Lee <christina@xanadu.ai>
Date:   Tue Jan 21 16:21:11 2025 -0500

    [Capture] Add backprop validation (#6852)

    **Context:**

    We currently use un-validated backprop for differentiation with program
    capture. This leads to some unintuitive errors if you try and take a
    gradient on lightning with capture enabled.

    **Description of the Change:**

    Adds some validation to make sure the device supports backprop. Adds the
    backprop logic to a `_backprop` jvp function, and dispatches to that
    method based on the diff method.

    **Benefits:**

    Improved error messages when backprop or the requested diff method isn't
    supported.

    **Possible Drawbacks:**

    The code currently is a little clunky, but it is private so we should be
    able to move things around once we have more information.

    **Related GitHub Issues:**

    [sc-82166]

commit 98bb29b
Author: Isaac De Vlugt <34751083+isaacdevlugt@users.noreply.github.com>
Date:   Tue Jan 21 14:14:39 2025 -0500

    `lie_closure_dense` typo in docstring (#6858)

    **Context:**

    **Description of the Change:** Docstring code example

    **Benefits:** Docstring code example works

    **Possible Drawbacks:** 0️⃣

    **Related GitHub Issues:**

commit fe9c9a1
Author: Yushao Chen (Jerry) <chenys13@outlook.com>
Date:   Tue Jan 21 13:48:22 2025 -0500

    Fix the deprecated usage of `MultiControlledX` in labs (#6862)

    **Context:**
    In the tests
    `pennylane/labs/tests/resource_estimation/ops/op_math/test_controlled_ops.py`
    there are several test suites that were using the deprecated interfaces
    of `MultiControlledX`, which will fail after the removal of
    corresponding deprecated `control_wires` arg.
    We fix the usages in this PR.

    **Description of the Change:**

    **Benefits:**

    **Possible Drawbacks:**

    **Related GitHub Issues:**

commit 872607d
Author: Mudit Pandey <mudit.pandey@xanadu.ai>
Date:   Tue Jan 21 11:48:25 2025 -0500

    Bump `torch` version in CI to 2.5.0 (#6868)

    As name says. We should be testing against the latest version of
    `torch`. This PR updates the torch version used in CI to `2.5.0`.

    Also updated torch installation to use `~=` instead of `==` for choosing
    the version, so that bug fix releases (highest possible `2.5.X`) are
    automatically used instead of sticking with `2.5.0`. This should be a
    safe change to make, as torch only adds bug fixes to patch releases and
    no breaking changes
    ([ref](https://lightning.ai/docs/pytorch/stable/versioning.html)).

commit 633b5bd
Author: Mudit Pandey <mudit.pandey@xanadu.ai>
Date:   Tue Jan 21 10:14:00 2025 -0500

    Update docs workflow schedule (#6867)

    [sc-82706]

    Docs workflow currently opens a PR to update stable deps on Wednesday,
    while the tests workflow opens the stable deps update PR on Monday.
    Having 2 PRs per week for the same thing is annoying, so the schedule
    time for the docs workflow should be the same as the tests workflow.

commit 37de975
Author: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Date:   Tue Jan 21 09:31:51 2025 -0500

    Update stable dependency files (#6856)

commit a97fca5
Author: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai>
Date:   Tue Jan 21 09:51:51 2025 +0000

    [no ci] bump nightly version

commit 5eaaccb
Author: David Wierichs <david.wierichs@xanadu.ai>
Date:   Tue Jan 21 10:05:20 2025 +0100

    Add some explanations on `NonInterpPrimitive` class (#6851)

    **Context:**
    The capture module uses a variant of `jax.core.Primitive` called
    `NonInterpPrimitive`.
    There were questions about why we need this and what it does.

    **Description of the Change:**
    This PR only adds some explanations to the respective `md` file to
    motivate our usage of this primitive variant.

    **Benefits:**
    Explain code

    **Possible Drawbacks:**
    N/A

    **Related GitHub Issues:**

    ---------

    Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
    Co-authored-by: Pietropaolo Frisoni <pietropaolo.frisoni@xanadu.ai>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@mudit2812 mudit2812 mudit2812 approved these changes

@andrijapau andrijapau andrijapau 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.
3 participants
@albi3ro @mudit2812 @andrijapau