-
Notifications
You must be signed in to change notification settings - Fork 598
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.
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
Move device preprocessing to inner execute #5791
Conversation
…o inner-execute
…o inner-execute
There was a problem hiding this 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, thanks @astralcai !
Just had a small suggestion for one test, and wanted to follow up on open comments.
Co-authored-by: David Wierichs <[email protected]>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @astralcai
…o inner-execute
tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
Show resolved
Hide resolved
Co-authored-by: Christina Lee <[email protected]>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🎉
**Context:** Currently blocked by PennyLaneAI/pennylane#5791. Device preprocessing is being moved inside the ML boundary, so the interface information needed for correct `hw-like` postselection is being lost when using `dynamic_one_shot`. **Description of the Change:** `qml.devices.preprocess.mid_circuit_measurements` accepts an additional `interface` kwarg, which is now being used in `LightningQubit.preprocess`. **Benefits:** Expected behaviour with `mcm_method="one-shot"`, `postselect_mode="hw-like"`, and `jax.jit`. **Possible Drawbacks:** **Related GitHub Issues:** --------- Co-authored-by: ringo-but-quantum <[email protected]> Co-authored-by: Ali Asadi <[email protected]>
**Context:** As name says. Gradient workflows no longer raise errors after the merge of #5791 , but their correctness is yet to be verified. **Description of the Change:** * Updated casting rules in `dynamic_one_shot`'s processing function for tensorflow. * For the changes to be fully integrated, the way that the interface is passed around when calling a QNode needed to be changed, so the following changes were made: * `QNode` has updated behaviour for how `mcm_config` is used during execution. In `QNode._execution_component`, a copy of `self.execute_kwargs["mcm_config"]` is the source of truth, and in `qml.execute`, `config.mcm_config` is the source of truth. * Added a private `pad-invalid-samples` `postselect_mode`. The `postselect_mode` is switched to this automatically in `qml.execute` if executing with jax and shots and `postselect_mode == "hw-like"`. This way we standardize how the MCM transforms determine if jax is being used. * Updates to `capture` module to accommodate the above changes. **Benefits:** * `dynamic_one_shot` doesn't cast to interfaces inside the ML boundary * `dynamic_one_shot` works with tensorflow * Expanded tests **Possible Drawbacks:** **Related GitHub Issues:** Fixes #5736, #5710 Duplicate of #5861 which was closed due to release branch merge stuff. --------- Co-authored-by: Jay Soni <[email protected]> Co-authored-by: Astral Cai <[email protected]> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Yushao Chen (Jerry) <[email protected]> Co-authored-by: Christina Lee <[email protected]> Co-authored-by: Thomas R. Bromley <[email protected]> Co-authored-by: soranjh <[email protected]> Co-authored-by: Pietropaolo Frisoni <[email protected]> Co-authored-by: Ahmed Darwish <[email protected]> Co-authored-by: Utkarsh <[email protected]> Co-authored-by: David Wierichs <[email protected]> Co-authored-by: Christina Lee <[email protected]> Co-authored-by: Mikhail Andrenkov <[email protected]> Co-authored-by: Diego <[email protected]> Co-authored-by: Josh Izaac <[email protected]> Co-authored-by: Diego <[email protected]> Co-authored-by: Vincent Michaud-Rioux <[email protected]> Co-authored-by: lillian542 <[email protected]> Co-authored-by: Jack Brown <[email protected]> Co-authored-by: Paul Finlay <[email protected]> Co-authored-by: David Ittah <[email protected]> Co-authored-by: Cristian Emiliano Godinez Ramirez <[email protected]> Co-authored-by: Vincent Michaud-Rioux <[email protected]>
Context:
If there is a gradient transform, we want device preprocessing to happen within the ml boundary, i.e., in
inner_execute
, to avoid applying the gradient transform on transformed tapes produced by device preprocessing.Description of the Change:
inner_transform
keyword toqml.execute
QNode._execute_component
andqml.execute
to add device preprocessing to the inner transform program if thegradient_fn
is a gradient transform.Related GitHub Issues:
[sc-59107]