StateMP accepts wires #4570

timmysilv · 2023-09-05T13:36:24Z

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), and the wire_order given to process_state is the wire order of the state being passed in (which happens here, at a point where we do not know about the device or its wires).

Changes for pad dispatching:

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
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)

This reverts commit 9dc6e36.

doc/releases/changelog-dev.md

mudit2812

Here's a rocket 🚀

pennylane/measurements/state.py

rmoyard

It looks good to me, but there is an issue with the TensorFlow padding. Happy to approve after it is fixed.

pennylane/math/single_dispatch.py

rmoyard

It looks good on my side, but can we double check that it does not break lightning for example?

codecov · 2023-09-05T20:06:32Z

Codecov Report

Patch coverage: 100.00% and no project coverage change.

Comparison is base (0843183) 99.65% compared to head (f8fb825) 99.65%.
Report is 1 commits behind head on master.

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

Files Changed	Coverage Δ
pennylane/math/single_dispatch.py	`99.46% <100.00%> (+<0.01%)`	⬆️
pennylane/measurements/state.py	`100.00% <100.00%> (ø)`

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

timmysilv · 2023-09-05T20:26:16Z

just added batching. @rmoyard this shouldn't break lightning (or any plugin for that matter) because no old device uses StateMP.process_state to compute the state. They all use QubitDevice.statistics which just returns device.state

rmoyard

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.

timmysilv added 3 commits September 1, 2023 15:27

Revert "remove wires from StateMP"

1c2dedb

This reverts commit 9dc6e36.

extend state functionality with wire_order

5c4d792

changelog

60cf837

timmysilv commented Sep 5, 2023

View reviewed changes

doc/releases/changelog-dev.md Outdated Show resolved Hide resolved

Update doc/releases/changelog-dev.md

562edbb

timmysilv requested review from mudit2812, rmoyard and lillian542 September 5, 2023 13:37

fix torch import; remove unneeded code

1a0faae

mudit2812 approved these changes Sep 5, 2023

View reviewed changes

pennylane/measurements/state.py Show resolved Hide resolved

rmoyard reviewed Sep 5, 2023

View reviewed changes

pennylane/math/single_dispatch.py Show resolved Hide resolved

timmysilv and others added 2 commits September 5, 2023 15:10

handle torch padding in state code

b3ed333

Merge branch 'master' into state-accepts-wires

3946b1c

rmoyard self-requested a review September 5, 2023 19:39

rmoyard reviewed Sep 5, 2023

View reviewed changes

support batching

f8fb825

rmoyard self-requested a review September 5, 2023 20:50

rmoyard approved these changes Sep 5, 2023

View reviewed changes

timmysilv merged commit 39e6577 into master Sep 5, 2023
39 checks passed

timmysilv deleted the state-accepts-wires branch September 5, 2023 21:27

timmysilv mentioned this pull request Sep 6, 2023

DQ2: Add device wires to measurement process when needed #4580

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

StateMP accepts wires #4570

StateMP accepts wires #4570

timmysilv commented Sep 5, 2023 •

edited

Loading

mudit2812 left a comment

rmoyard left a comment

rmoyard left a comment

codecov bot commented Sep 5, 2023 •

edited

Loading

timmysilv commented Sep 5, 2023

rmoyard left a comment

StateMP accepts wires #4570

StateMP accepts wires #4570

Conversation

timmysilv commented Sep 5, 2023 • edited Loading

mudit2812 left a comment

Choose a reason for hiding this comment

rmoyard left a comment

Choose a reason for hiding this comment

rmoyard left a comment

Choose a reason for hiding this comment

codecov bot commented Sep 5, 2023 • edited Loading

Codecov Report

timmysilv commented Sep 5, 2023

rmoyard left a comment

Choose a reason for hiding this comment

timmysilv commented Sep 5, 2023 •

edited

Loading

codecov bot commented Sep 5, 2023 •

edited

Loading