Merge branch 'master' into qcut-fix-expansion-obs-bug

doc/releases/changelog-dev.md

@@ -78,6 +78,9 @@
   configurations.
   [(#2169)](https://github.com/PennyLaneAI/pennylane/pull/2169)
 
+  The postprocessing function for the `cut_circuit` transform has been added.
+  [(#2192)](https://github.com/PennyLaneAI/pennylane/pull/2192)
+
 <h3>Improvements</h3>
 
 * The `gradients` module has been streamlined and special-purpose functions

pennylane/grouping/pauli.py

@@ -309,8 +309,11 @@ def partition_pauli_group(n_qubits: int) -> List[List[str]]:
     if not isinstance(n_qubits, int):
         raise TypeError("Must specify an integer number of qubits.")
 
-    if n_qubits <= 0:
-        raise ValueError("Number of qubits must be at least 1.")
+    if n_qubits < 0:
+        raise ValueError("Number of qubits must be at least 0.")
+
+    if n_qubits == 0:
+        return [[""]]
 
     strings = set()  # tracks all the strings that have already been grouped
     groups = []

pennylane/transforms/__init__.py

@@ -97,6 +97,7 @@
     ~transforms.graph_to_tape
     ~transforms.expand_fragment_tapes
     ~transforms.contract_tensors
+    ~transforms.qcut_processing_fn
 
 Transforms that act on tapes
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -177,4 +178,5 @@
     graph_to_tape,
     expand_fragment_tapes,
     contract_tensors,
+    qcut_processing_fn,
 )

pennylane/transforms/qcut.py

@@ -21,8 +21,9 @@
 from itertools import product
 from typing import List, Sequence, Tuple
 
-import pennylane as qml
 from networkx import MultiDiGraph, weakly_connected_components
+
+import pennylane as qml
 from pennylane import apply, expval
 from pennylane.grouping import string_to_pauli_word
 from pennylane.measure import MeasurementProcess
@@ -597,12 +598,12 @@ def contract_tensors(
             each tensor
         use_opt_einsum (bool): Determines whether to use the
             `opt_einsum <https://dgasmith.github.io/opt_einsum/>`__ package. This package is useful
-            for tensor contractions of large networks but must be installed separately using, e.g.,
-            ``pip install opt_einsum``. Both settings for ``use_opt_einsum`` result in a
+            for faster tensor contractions of large networks but must be installed separately using,
+            e.g., ``pip install opt_einsum``. Both settings for ``use_opt_einsum`` result in a
             differentiable contraction.
 
     Returns:
-        float or array-like: the result of contracting the tensor network
+        float or tensor_like: the result of contracting the tensor network
 
     **Example**
 
@@ -679,3 +680,166 @@ def contract_tensors(
     kwargs = {} if use_opt_einsum else {"like": tensors[0]}
 
     return contract(eqn, *tensors, **kwargs)
+
+
+CHANGE_OF_BASIS = qml.math.array(
+    [[1.0, 1.0, 0.0, 0.0], [-1.0, -1.0, 2.0, 0.0], [-1.0, -1.0, 0.0, 2.0], [1.0, -1.0, 0.0, 0.0]]
+)
+
+
+def _process_tensor(results, n_prep: int, n_meas: int):
+    """Convert a flat slice of an individual circuit fragment's execution results into a tensor.
+
+    This function performs the following steps:
+
+    1. Reshapes ``results`` into the intermediate shape ``(4,) * n_prep + (4**n_meas,)``
+    2. Shuffles the final axis to follow the standard product over measurement settings. E.g., for
+      ``n_meas = 2`` the standard product is: II, IX, IY, IZ, XI, ..., ZY, ZZ while the input order
+      will be the result of ``qml.grouping.partition_pauli_group(2)``, i.e., II, IZ, ZI, ZZ, ...,
+      YY.
+    3. Reshapes into the final target shape ``(4,) * (n_prep + n_meas)``
+    4. Performs a change of basis for the preparation indices (the first ``n_prep`` indices) from
+       the |0>, |1>, |+>, |+i> basis to the I, X, Y, Z basis using ``CHANGE_OF_BASIS``.
+
+    Args:
+        results (tensor_like): the input execution results
+        n_prep (int): the number of preparation nodes in the corresponding circuit fragment
+        n_meas (int): the number of measurement nodes in the corresponding circuit fragment
+
+    Returns:
+        tensor_like: the corresponding fragment tensor
+    """
+    n = n_prep + n_meas
+    dim_meas = 4**n_meas
+
+    # Step 1
+    intermediate_shape = (4,) * n_prep + (dim_meas,)
+    intermediate_tensor = qml.math.reshape(results, intermediate_shape)
+
+    # Step 2
+    grouped = qml.grouping.partition_pauli_group(n_meas)
+    grouped_flat = [term for group in grouped for term in group]
+    order = qml.math.argsort(grouped_flat)
+
+    if qml.math.get_interface(intermediate_tensor) == "tensorflow":
+        # TensorFlow does not support slicing
+        intermediate_tensor = qml.math.gather(intermediate_tensor, order, axis=-1)
+    else:
+        sl = [slice(None)] * n_prep + [order]
+        intermediate_tensor = intermediate_tensor[tuple(sl)]
+
+    # Step 3
+    final_shape = (4,) * n
+    final_tensor = qml.math.reshape(intermediate_tensor, final_shape)
+
+    # Step 4
+    change_of_basis = qml.math.convert_like(CHANGE_OF_BASIS, intermediate_tensor)
+
+    for i in range(n_prep):
+        axes = [[1], [i]]
+        final_tensor = qml.math.tensordot(change_of_basis, final_tensor, axes=axes)
+
+    axes = list(reversed(range(n_prep))) + list(range(n_prep, n))
+
+    # Use transpose to reorder indices. We must do this because tensordot returns a tensor whose
+    # indices are ordered according to the uncontracted indices of the first tensor, followed
+    # by the uncontracted indices of the second tensor. For example, calculating C_kj T_ij returns
+    # a tensor T'_ki rather than T'_ik.
+    final_tensor = qml.math.transpose(final_tensor, axes=axes)
+
+    final_tensor *= qml.math.power(2, -(n_meas + n_prep) / 2)
+    return final_tensor
+
+
+def _to_tensors(
+    results,
+    prepare_nodes: Sequence[Sequence[PrepareNode]],
+    measure_nodes: Sequence[Sequence[MeasureNode]],
+) -> List:
+    """Process a flat list of execution results from all circuit fragments into the corresponding
+    tensors.
+
+    This function slices ``results`` according to the expected size of fragment tensors derived from
+    the ``prepare_nodes`` and ``measure_nodes`` and then passes onto ``_process_tensor`` for further
+    transformation.
+
+    Args:
+        results (tensor_like): A collection of execution results, provided as a flat tensor,
+            corresponding to the expansion of circuit fragments in the communication graph over
+            measurement and preparation node configurations. These results are processed into
+            tensors by this function.
+        prepare_nodes (Sequence[Sequence[PrepareNode]]): a sequence whose length is equal to the
+            number of circuit fragments, with each element used here to determine the number of
+            preparation nodes in a given fragment
+        measure_nodes (Sequence[Sequence[MeasureNode]]): a sequence whose length is equal to the
+            number of circuit fragments, with each element used here to determine the number of
+            measurement nodes in a given fragment
+
+    Returns:
+        List[tensor_like]: the tensors for each circuit fragment in the communication graph
+    """
+    ctr = 0
+    tensors = []
+
+    for p, m in zip(prepare_nodes, measure_nodes):
+        n_prep = len(p)
+        n_meas = len(m)
+        n = n_prep + n_meas
+
+        dim = 4**n
+        results_slice = results[ctr : dim + ctr]
+
+        tensors.append(_process_tensor(results_slice, n_prep, n_meas))
+
+        ctr += dim
+
+    if len(results) != ctr:
+        raise ValueError(f"The results argument should be a flat list of length {ctr}")
+
+    return tensors
+
+
+def qcut_processing_fn(
+    results: Sequence[Sequence],
+    communication_graph: MultiDiGraph,
+    prepare_nodes: Sequence[Sequence[PrepareNode]],
+    measure_nodes: Sequence[Sequence[MeasureNode]],
+    use_opt_einsum: bool = False,
+):
+    """Processing function for the :func:`cut_circuit` transform.
+
+    .. note::
+
+        This function is designed for use as part of the circuit cutting workflow. Check out the
+        :doc:`transforms </code/qml_transforms>` page for more details.
+
+    Args:
+        results (Sequence[Sequence]): A collection of execution results corresponding to the
+            expansion of circuit fragments in the ``communication_graph`` over measurement and
+            preparation node configurations. These results are processed into tensors and then
+            contracted.
+        communication_graph (MultiDiGraph): the communication graph determining connectivity between
+            circuit fragments
+        prepare_nodes (Sequence[Sequence[PrepareNode]]): a sequence of size
+            ``len(communication_graph.nodes)`` that determines the order of preparation indices in
+            each tensor
+        measure_nodes (Sequence[Sequence[MeasureNode]]): a sequence of size
+            ``len(communication_graph.nodes)`` that determines the order of measurement indices in
+            each tensor
+        use_opt_einsum (bool): Determines whether to use the
+            `opt_einsum <https://dgasmith.github.io/opt_einsum/>`__ package. This package is useful
+            for faster tensor contractions of large networks but must be installed separately using,
+            e.g., ``pip install opt_einsum``. Both settings for ``use_opt_einsum`` result in a
+            differentiable contraction.
+
+    Returns:
+        float or tensor_like: the output of the original uncut circuit arising from contracting
+        the tensor network of circuit fragments
+    """
+    flat_results = qml.math.concatenate(results)
+
+    tensors = _to_tensors(flat_results, prepare_nodes, measure_nodes)
+    result = contract_tensors(
+        tensors, communication_graph, prepare_nodes, measure_nodes, use_opt_einsum
+    )
+    return result

tests/grouping/test_pauli_group.py

@@ -263,5 +263,9 @@ def test_invalid_input(self):
         with pytest.raises(TypeError, match="Must specify an integer number"):
             partition_pauli_group("3")
 
-        with pytest.raises(ValueError, match="Number of qubits must be at least 1"):
+        with pytest.raises(ValueError, match="Number of qubits must be at least 0"):
             partition_pauli_group(-1)
+
+    def test_zero(self):
+        """Test if [[""]] is returned with zero qubits"""
+        assert partition_pauli_group(0) == [[""]]