-
Notifications
You must be signed in to change notification settings - Fork 585
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
StateMP accepts wires #4570
StateMP accepts wires #4570
Conversation
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.
Here's a rocket 🚀
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.
It looks good to me, but there is an issue with the TensorFlow padding. Happy to approve after it is fixed.
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.
It looks good on my side, but can we double check that it does not break lightning for example?
Codecov ReportPatch coverage:
Additional details and impacted files@@ Coverage Diff @@
## master #4570 +/- ##
=======================================
Coverage 99.65% 99.65%
=======================================
Files 376 376
Lines 33108 33135 +27
=======================================
+ Hits 32995 33022 +27
Misses 113 113
☔ View full report in Codecov by Sentry. |
just added batching. @rmoyard this shouldn't break lightning (or any plugin for that matter) because no old device uses |
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, great changes 💯
**Context:** In the new device interface, a device need not have wires. In the case when it does provide wires, it would be best if measurement processes that don't require wires at all (`state()`, `sample()`, `probs()`) know to expand over the extra wires. To be able to do this, they need to save the wires somewhere, and `self.wires` kinda makes sense for this:) **Description of the Change:** - Allow initialization of a `StateMP` with wires (now that it doesn't indicate a density-matrix measurement). - When given a strict subset of wires, `process_state` pads the state with zeros for all other wire axes - When given a different `wire_order` from `self.wires`, `process_state` will transpose the state to match the MP wires **Note on wire ordering:** I use the convention that `StateMP.wires` is the global wire order we'd like to return when used (so we can swap out measurement processes [here](https://github.com/PennyLaneAI/pennylane/blob/078c4503b87e988f0ca68344ff57687d36f3e7fa/pennylane/devices/experimental/default_qubit_2.py#L217)), and the `wire_order` given to `process_state` is the wire order of the `state` being passed in (which happens [here](https://github.com/PennyLaneAI/pennylane/blob/078c4503b87e988f0ca68344ff57687d36f3e7fa/pennylane/devices/qubit/simulate.py#L91), at a point where we do not know about the device or its wires). **Changes for `pad` dispatching:** 1. I'm not sure what `qml.math.pad(x, like="torch")` was even calling, but the docs suggest to use `torch.nn.functional.pad` so I registered it 2. The tensorflow `pad` assumes padding over more than one dimension, so I "hacked" it to add an empty broadcast dimension if needed (in other words, turn `[pad]` into `[[pad]]`) **Benefits:** - `StateMP.process_state` is more powerful! - Allows me to open a follow-up PR where DQ2 expands results for the above listed MPs **Possible Drawbacks:** - Doesn't support batching (yet, I can add in this PR if necessary, just wanted something that worked) --------- Co-authored-by: Romain Moyard <rmoyard@gmail.com>
**Context:** Follow-up to #4570, DQ2 can now add wires to all measurement processes that might need them. **Description of the Change:** if DQ2 has wires set, and any measurements on a tape have no obs or wires, replace them with a duplicate that has wires set to `dev.wires`. This should only happen with `sample()`, `probs()`, `counts()` and `state()` afaik. New things: - `QuantumScript.op_wires`: The wires acted on by operations. Useful if some wires are only measured - `QuantumScript.map_to_standard_wires`: Maps to the wire-order that DQ2 expects. Docs with examples [here](https://xanaduai-pennylane--4580.com.readthedocs.build/en/4580/code/api/pennylane.tape.QuantumScript.html#pennylane.tape.QuantumScript.map_to_standard_wires) - `simulate.py::expand_over_state_wires()`: the same as `StateMP.process_state` but without the unflattening/flattening nonsense. I prefer this because we can keep it as a subroutine used by `get_final_state` as an optimization without affecting the rest of PL (no test changes needed! just added one for broadcasting coverage when wire mapping is needed) I also added `qml.Identity` instances to some measurements tests in the past to make up for this change not existing, so I'm removing them in this PR. **Benefits:** - Device wires are used as expected! - We can ignore extra wires only used in measurement as a nice optimization **Possible Drawbacks:** - Code is increasing in complexity... but I've kept everything isolated! - Might want to delete the code from StateMP to avoid duplication.
**Context:** Follow-up to #4570, DQ2 can now add wires to all measurement processes that might need them. **Description of the Change:** if DQ2 has wires set, and any measurements on a tape have no obs or wires, replace them with a duplicate that has wires set to `dev.wires`. This should only happen with `sample()`, `probs()`, `counts()` and `state()` afaik. New things: - `QuantumScript.op_wires`: The wires acted on by operations. Useful if some wires are only measured - `QuantumScript.map_to_standard_wires`: Maps to the wire-order that DQ2 expects. Docs with examples [here](https://xanaduai-pennylane--4580.com.readthedocs.build/en/4580/code/api/pennylane.tape.QuantumScript.html#pennylane.tape.QuantumScript.map_to_standard_wires) - `simulate.py::expand_over_state_wires()`: the same as `StateMP.process_state` but without the unflattening/flattening nonsense. I prefer this because we can keep it as a subroutine used by `get_final_state` as an optimization without affecting the rest of PL (no test changes needed! just added one for broadcasting coverage when wire mapping is needed) I also added `qml.Identity` instances to some measurements tests in the past to make up for this change not existing, so I'm removing them in this PR. **Benefits:** - Device wires are used as expected! - We can ignore extra wires only used in measurement as a nice optimization **Possible Drawbacks:** - Code is increasing in complexity... but I've kept everything isolated! - Might want to delete the code from StateMP to avoid duplication.
Context:
In the new device interface, a device need not have wires. In the case when it does provide wires, it would be best if measurement processes that don't require wires at all (
state()
,sample()
,probs()
) know to expand over the extra wires. To be able to do this, they need to save the wires somewhere, andself.wires
kinda makes sense for this:)Description of the Change:
StateMP
with wires (now that it doesn't indicate a density-matrix measurement).process_state
pads the state with zeros for all other wire axeswire_order
fromself.wires
,process_state
will transpose the state to match the MP wiresNote on wire ordering:
I use the convention that
StateMP.wires
is the global wire order we'd like to return when used (so we can swap out measurement processes here), and thewire_order
given toprocess_state
is the wire order of thestate
being passed in (which happens here, at a point where we do not know about the device or its wires).Changes for
pad
dispatching:qml.math.pad(x, like="torch")
was even calling, but the docs suggest to usetorch.nn.functional.pad
so I registered itpad
assumes padding over more than one dimension, so I "hacked" it to add an empty broadcast dimension if needed (in other words, turn[pad]
into[[pad]]
)Benefits:
StateMP.process_state
is more powerful!Possible Drawbacks: