Add examples with mid-circuit measurements across the documentation (#…

…5485)

**Context:**
We want to make mid-circuit measurements a first-class member operation
of PennyLane.
Accordingly, they should be featured in examples across the
documentation, raising visibility and showcasing compatibilities between
advanced features and MCMs.

**Description of the Change:**
This PR adds MCMs
- to the main example circuit used in the docstring of `draw_mpl`,
highlighting the beautiful MCM drawing capabilities added last year.
- to the docstring of `classical_fisher`, demonstrating compatibility
across modules
- to the docstring of `QFT`, showcasing a nice usage example of MCMs in
the semiclassical Fourier transform. (Here one could argue that the
example is not _on_ `qml.QFT` but on a scenario where `qml.QFT` is used.

The PR also fixes some minor typos.

**Benefits:**
More visibility for MCMs and dynamic circuits in the docs.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
N/A

---------

Co-authored-by: Isaac De Vlugt <34751083+isaacdevlugt@users.noreply.github.com>
Co-authored-by: Isaac De Vlugt <isaacdevlugt@gmail.com>
Co-authored-by: Thomas R. Bromley <49409390+trbromley@users.noreply.github.com>

doc/introduction/compiling_circuits.rst

@@ -43,17 +43,17 @@ RX(0.09999999999999964, wires=[0])
 RX(11.336370614359172, wires=[0])
 >>> qml.simplify(qml.ops.Pow(qml.RX(1, 0), 3))
 RX(3.0, wires=[0])
->>> qml.simplify(qml.sum(qml.PauliY(3), qml.PauliY(3)))
+>>> qml.simplify(qml.sum(qml.Y(3), qml.Y(3)))
 2 * Y(3)
 >>> qml.simplify(qml.RX(1, 0) @ qml.RX(1, 0))
 RX(2.0, wires=[0])
->>> qml.simplify(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
+>>> qml.simplify(qml.prod(qml.X(0), qml.Z(0)))
 -1j * Y(0)
 
 Now lets simplify a nested operator:
 
->>> sum_op = qml.RX(1, 0) + qml.PauliX(0)
->>> prod1 = qml.PauliX(0) @ sum_op
+>>> sum_op = qml.RX(1, 0) + qml.X(0)
+>>> prod1 = qml.X(0) @ sum_op
 >>> nested_op = prod1 @ qml.RX(1, 0)
 >>> qml.simplify(nested_op)
 (X(0) @ RX(2.0, wires=[0])) + RX(1.0, wires=[0])
@@ -62,19 +62,19 @@ Several simplifications steps are happening here. First of all, the nested produ
 
 .. code-block:: python
 
-    qml.prod(qml.PauliX(0), qml.sum(qml.RX(1, 0), qml.PauliX(0)), qml.RX(1, 0))
+    qml.prod(qml.X(0), qml.sum(qml.RX(1, 0), qml.X(0)), qml.RX(1, 0))
 
 Then the product of sums is transformed into a sum of products:
 
 .. code-block:: python
 
-    qml.sum(qml.prod(qml.PauliX(0), qml.RX(1, 0), qml.RX(1, 0)), qml.prod(qml.PauliX(0), qml.PauliX(0), qml.RX(1, 0)))
+    qml.sum(qml.prod(qml.X(0), qml.RX(1, 0), qml.RX(1, 0)), qml.prod(qml.X(0), qml.X(0), qml.RX(1, 0)))
 
 And finally like terms in the obtained products are grouped together, removing all identities: 
 
 .. code-block:: python
 
-    qml.sum(qml.prod(qml.PauliX(0), qml.RX(2, 0)), qml.RX(1, 0))
+    qml.sum(qml.prod(qml.X(0), qml.RX(2, 0)), qml.RX(1, 0))
 
 As mentioned earlier we can also simplify QNode objects to, for example, group rotation gates:
 
@@ -219,7 +219,7 @@ For example, suppose we would like to implement the following QNode:
 
     def circuit(weights):
         qml.BasicEntanglerLayers(weights, wires=[0, 1, 2])
-        return qml.expval(qml.PauliZ(0))
+        return qml.expval(qml.Z(0))
 
     original_dev = qml.device("default.qubit", wires=3)
     original_qnode = qml.QNode(circuit, original_dev)
@@ -333,12 +333,13 @@ by devices to make such measurements possible.
 On a lower level, the :func:`~.pennylane.pauli.group_observables` function can be used to split lists of
 observables and coefficients:
 
->>> obs = [qml.PauliY(0), qml.PauliX(0) @ qml.PauliX(1), qml.PauliZ(1)]
+>>> obs = [qml.Y(0), qml.X(0) @ qml.X(1), qml.Z(1)]
 >>> coeffs = [1.43, 4.21, 0.97]
->>> obs_groupings, coeffs_groupings = qml.pauli.group_observables(obs, coeffs, 'anticommuting', 'lf')
+>>> groupings = qml.pauli.group_observables(obs, coeffs, 'anticommuting', 'lf')
+>>> obs_groupings, coeffs_groupings = groupings
 >>> obs_groupings
 [[Z(1), X(0) @ X(1)], [Y(0)]]
 >>> coeffs_groupings
 [[0.97, 4.21], [1.43]]
 
-This and more logic to manipulate Pauli observables is found in the :doc:`pauli module <../code/qml_pauli>`.
+This and more logic to manipulate Pauli observables is found in the :doc:`pauli module <../code/qml_pauli>`.

doc/introduction/compiling_workflows.rst

@@ -37,7 +37,7 @@ Check out the Catalyst documentation for
 Just-in-time compilation
 ------------------------
 
-Using Catalyst with PennyLane is a simple as using the :func:`@qjit <.qjit>` decorator to
+Using Catalyst with PennyLane is as simple as using the :func:`@qjit <.qjit>` decorator to
 compile your hybrid workflows:
 
 .. code-block:: python
@@ -53,7 +53,7 @@ compile your hybrid workflows:
         qml.RX(jnp.sin(params[0]) ** 2, wires=1)
         qml.CRY(params[0], wires=[0, 1])
         qml.RX(jnp.sqrt(params[1]), wires=1)
-        return qml.expval(qml.PauliZ(1))
+        return qml.expval(qml.Z(1))
 
 The :func:`~.qjit` decorator can also be used on hybrid functions --
 that is, functions that include both QNodes and classical processing.
@@ -78,7 +78,7 @@ using ``@jax.jit``:
         # initial parameter
         params = jnp.array([0.54, 0.3154])
 
-        # define the optimizer
+        # define the optimizer using a qjit-decorated function
         opt = jaxopt.GradientDescent(circuit, stepsize=0.4)
         update = lambda i, args: tuple(opt.update(*args))
 
@@ -111,7 +111,7 @@ rather than in Python at compile time. You can enable this feature via the
         else:
             qml.T(wires=0)
 
-        return qml.expval(qml.PauliZ(0))
+        return qml.expval(qml.Z(0))
 
 >>> circuit(3)
 array(0.)
@@ -149,7 +149,7 @@ decorator:
         qml.RX(x[0], wires=0)
         qml.RY(x[1], wires=1)
         qml.CNOT(wires=[0, 1])
-        return qml.expval(qml.PauliY(0))
+        return qml.expval(qml.Y(0))
 
 >>> circuit(jnp.array([0.5, 1.4]))
 -0.47244976756708373

doc/introduction/inspecting_circuits.rst

@@ -47,7 +47,7 @@ For example:
         qml.Toffoli(wires=(0, 1, 2))
         qml.CRY(x[1], wires=(0, 1))
         qml.Rot(x[2], x[3], y, wires=0)
-        return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1))
+        return qml.expval(qml.Z(0)), qml.expval(qml.X(1))
 
 
 We can now use the :func:`~pennylane.specs` transform to generate a function that returns
@@ -85,17 +85,20 @@ For example:
 
 .. code-block:: python
 
-    dev = qml.device('lightning.qubit', wires=(0,1,2,3))
+    dev = qml.device('default.qubit')
 
     @qml.qnode(dev)
     def circuit(x, z):
         qml.QFT(wires=(0,1,2,3))
         qml.IsingXX(1.234, wires=(0,2))
         qml.Toffoli(wires=(0,1,2))
+        mcm = qml.measure(1)
+        mcm_out = qml.measure(2)
         qml.CSWAP(wires=(0,2,3))
         qml.RX(x, wires=0)
+        qml.cond(mcm, qml.RY)(np.pi / 4, wires=3)
         qml.CRZ(z, wires=(3,0))
-        return qml.expval(qml.PauliZ(0))
+        return qml.expval(qml.Z(0)), qml.probs(op=mcm_out)
 
 
     fig, ax = qml.draw_mpl(circuit)(1.2345,1.2345)
@@ -107,10 +110,12 @@ For example:
     :target: javascript:void(0);
 
 >>> print(qml.draw(circuit)(1.2345,1.2345))
-0: ─╭QFT─╭IsingXX(1.23)─╭●─╭●─────RX(1.23)─╭RZ(1.23)─┤  <Z>
-1: ─├QFT─│──────────────├●─│───────────────│─────────┤     
-2: ─├QFT─╰IsingXX(1.23)─╰X─├SWAP───────────│─────────┤     
-3: ─╰QFT───────────────────╰SWAP───────────╰●────────┤     
+0: ─╭QFT─╭IsingXX(1.23)─╭●───────────╭●─────RX(1.23)─╭RZ(1.23)─┤  <Z>
+1: ─├QFT─│──────────────├●──┤↗├──────│───────────────│─────────┤
+2: ─├QFT─╰IsingXX(1.23)─╰X───║───┤↗├─├SWAP───────────│─────────┤
+3: ─╰QFT─────────────────────║────║──╰SWAP──RY(0.79)─╰●────────┤
+                             ╚════║═════════╝
+                                  ╚════════════════════════════╡  Probs[MCM]
 
 More information, including various fine-tuning options, can be found in
 the :doc:`drawing module <../code/qml_drawer>`.
@@ -128,20 +133,25 @@ Currently supported devices include:
 * ``default.mixed``: each snapshot saves the density matrix
 * ``default.gaussian``: each snapshot saves the covariance matrix and vector of means
 
-During normal execution, the snapshots are ignored:
+A :class:`~pennylane.Snapshot` can be used in a QNode like any other operation:
 
 .. code-block:: python
 
     dev = qml.device("default.qubit", wires=2)
 
     @qml.qnode(dev, interface=None)
     def circuit():
-        qml.Snapshot(measurement=qml.expval(qml.PauliZ(0)))
+        qml.Snapshot(measurement=qml.expval(qml.Z(0)))
         qml.Hadamard(wires=0)
         qml.Snapshot("very_important_state")
         qml.CNOT(wires=[0, 1])
         qml.Snapshot()
-        return qml.expval(qml.PauliX(0))
+        return qml.expval(qml.X(0))
+
+During normal execution, the snapshots are ignored:
+
+>>> circuit()
+0.0
 
 However, when using the :func:`~pennylane.snapshots`
 transform, intermediate device states will be stored and returned alongside the
@@ -153,6 +163,9 @@ results.
 2: array([0.707+0.j, 0.+0.j, 0.+0.j, 0.707+0.j]),
 'execution_results': 0.0}
 
+All snapshots are numbered with consecutive integers, and if no tag was provided,
+the number of a snapshot is used as a key in the output dictionary instead.
+
 Graph representation
 --------------------
 
@@ -185,7 +198,7 @@ or to check whether two gates causally influence each other.
         qml.CNOT([1, 2])
         qml.CNOT([2, 3])
         qml.CNOT([3, 1])
-        return qml.expval(qml.PauliZ(0))
+        return qml.expval(qml.Z(0))
 
 
     circuit()
@@ -252,7 +265,7 @@ pairwise commutation:
 ...     qml.Hadamard(wires=2)
 ...     qml.CRZ(z, wires=[2, 0])
 ...     qml.RY(-y, wires=1)
-...     return qml.expval(qml.PauliZ(0))
+...     return qml.expval(qml.Z(0))
 >>> dag_fn = qml.commutation_dag(circuit)
 >>> dag = dag_fn(np.pi / 4, np.pi / 3, np.pi / 2)
 

pennylane/drawer/draw.py

@@ -370,10 +370,13 @@ def circuit(x, z):
             qml.QFT(wires=(0,1,2,3))
             qml.IsingXX(1.234, wires=(0,2))
             qml.Toffoli(wires=(0,1,2))
+            mcm = qml.measure(1)
+            mcm_out = qml.measure(2)
             qml.CSWAP(wires=(0,2,3))
             qml.RX(x, wires=0)
+            qml.cond(mcm, qml.RY)(np.pi / 4, wires=3)
             qml.CRZ(z, wires=(3,0))
-            return qml.expval(qml.Z(0))
+            return qml.expval(qml.Z(0)), qml.probs(op=mcm_out)
 
 
         fig, ax = qml.draw_mpl(circuit)(1.2345,1.2345)
@@ -449,8 +452,12 @@ def circuit2(x, y):
             box1 = plt.Rectangle((-0.5, -0.5), width=3.0, height=4.0, **options)
             ax.add_patch(box1)
 
-            ax.annotate("CSWAP", xy=(3, 2.5), xycoords='data', xytext=(3.8,1.5), textcoords='data',
+            ax.annotate("CSWAP", xy=(5, 2.5), xycoords='data', xytext=(5.8,1.5), textcoords='data',
                         arrowprops={'facecolor': 'black'}, fontsize=14)
+
+            ax.annotate("classical control flow", xy=(3.5, 4.2), xycoords='data', xytext=(0.8,4.2),
+                        textcoords='data', arrowprops={'facecolor': 'blue'}, fontsize=14,
+                        va="center")
             fig.show()
 
         .. figure:: ../../_static/draw_mpl/postprocessing.png
@@ -490,7 +497,7 @@ def circuit2(x, y):
             plt.rcParams['patch.linewidth'] = 4
             plt.rcParams['patch.force_edgecolor'] = True
             plt.rcParams['lines.color'] = 'indigo'
-            plt.rcParams['lines.linewidth'] = 5
+            plt.rcParams['lines.linewidth'] = 2
             plt.rcParams['figure.facecolor'] = 'ghostwhite'
 
             fig, ax = qml.draw_mpl(circuit, style="rcParams")(1.2345,1.2345)

pennylane/measurements/__init__.py

@@ -103,8 +103,10 @@ def circ(x, y):
 reset to the :math:`|0 \rangle` state by setting the ``reset`` keyword argument of ``qml.measure`` to ``True``.
 
 Users can also collect statistics on mid-circuit measurements along with other terminal measurements. Currently,
-``qml.expval``, ``qml.probs``, ``qml.sample``, ``qml.counts``, and ``qml.var`` are supported. Users have the
-ability to collect statistics on single measurement values.
+``qml.expval``, ``qml.probs``, ``qml.sample``, ``qml.counts``, and ``qml.var`` are supported. ``qml.probs``,
+``qml.sample``, and ``qml.counts`` support sequences of measurement values, ``qml.expval`` and ``qml.var`` do not.
+Statistics of arithmetic combinations of measurement values are supported by all but ``qml.probs``, and only as
+long as they are not collected in a sequence, e.g., ``[m1 + m2, m1 - m2]`` is not supported.
 
 .. code-block:: python
 

pennylane/qinfo/transforms.py

@@ -57,6 +57,23 @@ def circuit(x):
     tensor([[0.5+0.j, 0. +0.j],
             [0. +0.j, 0.5+0.j]], requires_grad=True)
 
+    This is equivalent to the state of the wire ``0`` after measuring the wire ``1``:
+
+    .. code-block:: python
+
+        @qml.qnode(dev)
+        def measured_circuit(x):
+            qml.IsingXX(x, wires=[0,1])
+            m = qml.measure(1)
+            return qml.density_matrix(wires=[0]), qml.probs(op=m)
+
+    >>> dm, probs = measured_circuit(np.pi/2)
+    >>> dm
+    tensor([[0.5+0.j, 0. +0.j],
+            [0. +0.j, 0.5+0.j]], requires_grad=True)
+    >>> probs
+    tensor([0.5, 0.5], requires_grad=True)
+
     .. seealso:: :func:`pennylane.density_matrix` and :func:`pennylane.math.reduce_dm`
     """
     # device_wires is provided by the custom QNode transform
@@ -495,30 +512,31 @@ def classical_fisher(qnode, argnums=0):
     .. code-block:: python
 
         import pennylane.numpy as pnp
-        n_wires = 2
 
-        dev = qml.device("default.qubit", wires=n_wires)
+        dev = qml.device("default.qubit")
 
         @qml.qnode(dev)
         def circ(params):
             qml.RX(params[0], wires=0)
-            qml.RX(params[1], wires=0)
-            qml.CNOT(wires=(0,1))
-            return qml.probs(wires=range(n_wires))
+            qml.CNOT([0, 1])
+            qml.CRY(params[1], wires=[1, 0])
+            qml.Hadamard(1)
+            return qml.probs(wires=[0, 1])
 
-    Executing this circuit yields the ``2**n_wires`` elements of :math:`p_\ell(\bm{\theta})`
+    Executing this circuit yields the ``2**2=4`` elements of :math:`p_\ell(\bm{\theta})`
 
+    >>> pnp.random.seed(25)
     >>> params = pnp.random.random(2)
     >>> circ(params)
-    [0.61281668 0.         0.         0.38718332]
+    [0.41850088 0.41850088 0.08149912 0.08149912]
 
     We can obtain its ``(2, 2)`` classical fisher information matrix (CFIM) by simply calling the function returned
     by ``classical_fisher()``:
 
     >>> cfim_func = qml.qinfo.classical_fisher(circ)
     >>> cfim_func(params)
-    [[1. 1.]
-     [1. 1.]]
+    [[ 0.901561 -0.125558]
+     [-0.125558  0.017486]]
 
     This function has the same signature as the :class:`.QNode`. Here is a small example with multiple arguments:
 
@@ -785,7 +803,7 @@ def fidelity(qnode0, qnode1, wires0, wires1):
     fidelity is simply
 
     .. math::
-        F( \ket{\psi} , \ket{\phi}) = \left|\braket{\psi, \phi}\right|^2
+        F( \ket{\psi} , \ket{\phi}) = \left|\braket{\psi| \phi}\right|^2
 
     .. note::
         The second state is coerced to the type and dtype of the first state. The fidelity is returned in the type

pennylane/shadows/__init__.py

@@ -74,24 +74,28 @@
 
 .. code-block:: python3
 
-    H = qml.Hamiltonian([1., 1.], [qml.Z(0) @ qml.Z(1), qml.X(0) @ qml.X(1)])
+    H = qml.Hamiltonian([1., 1.], [qml.Z(0) @ qml.Z(1), qml.X(0) @ qml.Z(1)])
 
-    dev = qml.device("default.qubit", wires=range(2), shots=10000)
+    dev = qml.device("default.qubit", shots=10000)
+
+    # shadow_expval + mid-circuit measurements require to defer measurements
+    @qml.defer_measurements
     @qml.qnode(dev)
     def qnode(x):
         qml.Hadamard(0)
         qml.CNOT((0,1))
         qml.RX(x, wires=0)
+        qml.measure(1)
         return qml.shadow_expval(H)
 
     x = np.array(0.5, requires_grad=True)
 
 The big advantage of this way of computing expectation values is that it is differentiable.
 
 >>> qnode(x)
-array(1.9242)
+array(0.8406)
 >>> qml.grad(qnode)(x)
--0.44999999999999984
+-0.49680000000000013
 
 There are more options for post-processing classical shadows in :class:`ClassicalShadow`.
 """