Allow postselecting on mid-circuit measurements #4604
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
trbromley
reviewed
Oct 19, 2023
timmysilv
reviewed
Oct 19, 2023
albi3ro
approved these changes
Oct 20, 2023
albi3ro
approved these changes
Oct 20, 2023
timmysilv
reviewed
Oct 20, 2023
timmysilv
approved these changes
Oct 20, 2023
mudit2812
added
the
merge-ready ✔️
mudit2812
added a commit
that referenced
this pull request
Nov 24, 2023
**Context:**
Adding support for `qml.cond` for mid-circuit measurement drawing with
`qml.draw`. This PR adds support for using one mid-circuit measurement
with multiple `qml.cond`s, where each `qml.cond` uses the **same**
mid-circuit measurement. An example of the drawing is shown below:
```python
import pennylane as qml
def circuit():
qml.RX(0.5, 0)
qml.RX(0.5, 1)
m0 = qml.measure(0, reset=True, postselect=1)
qml.cond(m0, qml.MultiControlledX)(wires=[1, 2, 3, 0], control_values="110")
qml.CNOT([3, 2])
qml.cond(m0, qml.ctrl(qml.MultiRZ, control=[1, 2], control_values=[True, False]))(0.5, wires=[0, 3])
return qml.expval(qml.PauliZ(0))
print(qml.draw(circuit)())
```
```
0: ──RX(0.50)──┤↗₁│ │0⟩─╭X────╭MultiRZ(0.50)─┤ <Z>
1: ──RX(0.50)───║────────├●────├●─────────────┤
2: ─────────────║────────├●─╭X─├○─────────────┤
3: ─────────────║────────╰○─╰●─╰MultiRZ(0.50)─┤
╚═════════╩═════╝
```
**Note:** This PR contains work that is foundational for follow up PRs
but isn't used for drawing the cases mentioned above. I'll add arbitrary
testing to get full coverage, but the aforementioned cases will be the
only ones considered to currently be complete.
**Description of the Change:**
`drawable_layers`:
* `drawable_layers` no longer stacks mid-circuit measurements used by a
`qml.cond` in the same layer.
* `drawable layers` now considers all wires under a mid-circuit
measurement used for conditioning to be used.
* `drawable_layers` now considers all wires under a `qml.cond` to be
used.
`tape_text`:
* Added helper functions to `tape_text.py` for correctly adding labels
for mid-circuit measurements and `qml.cond`s, as well as for correctly
padding the width of the layers for unused wires.
* Added helper function to `tape_text.py` for getting information about
connectivity between mid-circuit measurements and conditional operations
to keep track of when to start and stop drawing classical bits.
* `tape_text()` now tracks a `bit_map`, which stores a mapping between
mid-circuit measurements and classical bits that are drawn for them.
Classical bits are drawn using double lines `═`, and are placed under
the wires in the drawing. Each mid-circuit measurement used for
conditioning uses a classical bit. Measurements not used for
conditioning are not allocated a bit.
* `tape_text()` now pads the classical bits between each layer of the
drawing along with the wires.
* This PR also addresses #4809. The changes pertaining to that are
contained in `tape_text.py::_add_measurements()`.
Other changes:
* Added `label` to `qml.transforms.Conditional`.
* Adjusted `qml.equal` to consider postselection when comparing
`MidMeasureMP`s. This was missed in #4604.
**Benefits:**
Drawing support for mid-circuit measurements is better :)
**Possible Drawbacks:**
Code in `tape_text.py` is less readable. After a discussion with
@trbromley , we decided that the best course of action is to leave
refactoring as technical debt for the end of the release cycle. Reviews
pertaining to major refactors are still appreciated :)
**Related GitHub Issues:**
#4809
---------
Co-authored-by: Christina Lee <christina@xanadu.ai>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Context:
This PR adds support for postselecting on mid-circuit measurements with
default.qubit. The various changes to cover the edge cases of postselection are described below.Description of the Change:
Changes to measurements, math, and transforms:
postselectkwarg toqml.measure.defer_measurementsto support postselection:qml.Projectoras an operation when postselection is requested.broadcast_expandwhen postselection is requested and batch_size is notNone. This is done because there is a possibility of ending up withNaNvalues in the state, and we only want the state vector of the "bad" parameter to haveNaNvalues, not the entire state vector.default.qubit.qml.mathforisnan,full, andany. Why must tensorflow name their functions slightly differently than everyone else? The functions I needed to add dispatches for are calledtf.math.is_nan,tf.fill, andtf.reduce_anyrespectively 😢Changes to
default.qubit:qml.Projectorand renormalizing state insimulate.py:np.random.binomialto update the number of shots on the tape such that only a subset of the shots are then used. This is done to simulate behaviour on real hardware, where invalid samples would be discarded._FlexShotsclass which inheritsShotstosimulate.pythat allows shot vectors to contain 0 shots. We need this in cases when the probability of measuring the postselected state is close to 0 but not 0. In these cases, it is possible to end up with 0 samples that are valid, and we need to be able to handle that.sampling.pyfor catching errors raised by the rng if attempting to generate samples from nan probabilities. When this happens, samples of shape (shots, len(wires))with NaN values are returned.The following is done to measurements in cases where the state is invalid or the number of shots is 0:
Measurement processes return
NaNvalues within their expected shape. Eg,qml.probsreturns a2**nsize array with allNaNs, andqml.expvalreturns a singleNaN.QInfo measurements don't work for analytic execution, and
probsandcountsdon't work for finite shots if state hasNaNvalues.Other changes:
defer_measurementsin QNode construction if the device isdefault.qubit. This is done so that multiple tapes are not created by the transform in case there is postselection and broadcasting.TODOs:
tests.devices.qubit.test_sampling.TestInvalidStateSamples.Benefits:
Postselection is available on
default.qubit.Possible Drawbacks:
defer_measurementsfrom qnode construction if the new device was being used as postselection would split tapes if the tape was broadcasted. This causes the drawer to not work withdefault.qubitfor circuits with mid-circuit measurements unlessdefer_measurementsis explicitly applied to the QNode.Related GitHub Issues:
Here's a gist with the notebook I've been using to run examples locally. It has flags to test analytic vs finite shots, integer shots vs shot vector, exactly zero vs close to zero postselected state probability, and unit vs integration execution