ZNE local folding

doc/releases/changelog-dev.md

@@ -26,6 +26,40 @@
   Array([[1], [0], [1], [1], [0], [1],[0]], dtype=int64))
   ```
 
+* Zero-Noise Extrapolation (ZNE) local folding: Introduces the option to fold gates locally as well as the existing method of globally. Global folding (as in previous versions) applies the scale factor by forming the inverse of the entire quantum circuit (without measurements) and repeating the circuit with its inverse; local folding inserts per-gate folding sequences directly in place of each gate in the original circuit instead of applying the scale factor to the entire circuit at once. [(#1006)](https://github.com/PennyLaneAI/catalyst/pull/1006)
+
+  For example,
+
+  ```python
+  import jax
+  import pennylane as qml
+  from catalyst import qjit, mitigate_with_zne
+  from pennylane.transforms import exponential_extrapolate
+
+  dev = qml.device("lightning.qubit", wires=4, shots=5)
+
+  @qml.qnode(dev)
+  def circuit():
+    qml.Hadamard(wires=0)
+    qml.CNOT(wires=[0, 1])
+    return qml.expval(qml.PauliY(wires=0))
+
+  @qjit(keep_intermediate=True)
+  def mitigated_circuit():
+    s = jax.numpy.array([1, 2, 3])
+    return mitigate_with_zne(
+      circuit,
+      scale_factors=s,
+      extrapolate=exponential_extrapolate,
+      folding="all" # "all" for local or "global" for the original method (default being "global")
+    )()
+  ```
+
+  ```pycon
+  >>> circuit()
+  >>> mitigated_circuit()
+  ```
+
 <h3>Improvements</h3>
 
 <h3>Breaking changes</h3>
@@ -56,3 +90,4 @@ Romain Moyard,
 Erick Ochoa Lopez,
 Paul Haochen Wang,
 Sengthai Heng,
+Daniel Strano

frontend/catalyst/api_extensions/error_mitigation.py

@@ -56,6 +56,7 @@ def mitigate_with_zne(
             function.
         folding (str): Unitary folding technique to be used to scale the circuit. Possible values:
             - global: the global unitary of the input circuit is folded
+            - all: per-gate folding sequences replace original gates in-place in the circuit
 
     Returns:
         Callable: A callable object that computes the mitigated of the wrapped :class:`~.QNode`
@@ -117,6 +118,7 @@ def workflow(weights, s):
     >>> workflow(weights, scale_factors)
     Array(-0.19946598, dtype=float64)
     """
+
     kwargs = copy.copy(locals())
     kwargs.pop("fn")
 
@@ -175,7 +177,7 @@ def __call__(self, *args, **kwargs):
         except ValueError as e:
             raise ValueError(f"Folding type must be one of {list(map(str, Folding))}") from e
         # TODO: remove the following check once #755 is completed
-        if folding != Folding.GLOBAL:
+        if folding == Folding.RANDOM:
             raise NotImplementedError(f"Folding type {folding.value} is being developed")
 
         results = zne_p.bind(

frontend/test/pytest/test_mitigation.py

@@ -34,7 +34,8 @@ def skip_if_exponential_extrapolation_unstable(circuit_param, extrapolation_func
 
 @pytest.mark.parametrize("params", [0.1, 0.2, 0.3, 0.4, 0.5])
 @pytest.mark.parametrize("extrapolation", [quadratic_extrapolation, exponential_extrapolate])
-def test_single_measurement(params, extrapolation):
+@pytest.mark.parametrize("folding", ["global", "all"])
+def test_single_measurement(params, extrapolation, folding):
     """Test that without noise the same results are returned for single measurements."""
     skip_if_exponential_extrapolation_unstable(params, extrapolation)
 
@@ -52,15 +53,19 @@ def circuit(x):
     @catalyst.qjit
     def mitigated_qnode(args):
         return catalyst.mitigate_with_zne(
-            circuit, scale_factors=jax.numpy.array([1, 2, 3]), extrapolate=extrapolation
+            circuit,
+            scale_factors=jax.numpy.array([1, 2, 3]),
+            extrapolate=extrapolation,
+            folding=folding,
         )(args)
 
     assert np.allclose(mitigated_qnode(params), circuit(params))
 
 
 @pytest.mark.parametrize("params", [0.1, 0.2, 0.3, 0.4, 0.5])
 @pytest.mark.parametrize("extrapolation", [quadratic_extrapolation, exponential_extrapolate])
-def test_multiple_measurements(params, extrapolation):
+@pytest.mark.parametrize("folding", ["global", "all"])
+def test_multiple_measurements(params, extrapolation, folding):
     """Test that without noise the same results are returned for multiple measurements"""
     skip_if_exponential_extrapolation_unstable(params, extrapolation)
 
@@ -78,14 +83,18 @@ def circuit(x):
     @catalyst.qjit
     def mitigated_qnode(args):
         return catalyst.mitigate_with_zne(
-            circuit, scale_factors=jax.numpy.array([1, 2, 3]), extrapolate=extrapolation
+            circuit,
+            scale_factors=jax.numpy.array([1, 2, 3]),
+            extrapolate=extrapolation,
+            folding=folding,
         )(args)
 
     assert np.allclose(mitigated_qnode(params), circuit(params))
 
 
 @pytest.mark.parametrize("params", [0.1, 0.2, 0.3, 0.4, 0.5])
-def test_single_measurement_control_flow(params):
+@pytest.mark.parametrize("folding", ["global", "all"])
+def test_single_measurement_control_flow(params, folding):
     """Test that without noise the same results are returned for single measurement and with
     control flow."""
     dev = qml.device("lightning.qubit", wires=2)
@@ -113,9 +122,9 @@ def loop_1(i):  # pylint: disable=unused-argument
 
     @catalyst.qjit
     def mitigated_qnode(args, n):
-        return catalyst.mitigate_with_zne(circuit, scale_factors=jax.numpy.array([1, 2, 3]))(
-            args, n
-        )
+        return catalyst.mitigate_with_zne(
+            circuit, scale_factors=jax.numpy.array([1, 2, 3]), folding=folding
+        )(args, n)
 
     assert np.allclose(mitigated_qnode(params, 3), catalyst.qjit(circuit)(params, 3))
 
@@ -238,15 +247,16 @@ def circuit():
         return 0.0
 
     def mitigated_qnode():
-        return catalyst.mitigate_with_zne(circuit, scale_factors=[], folding="all")()
+        return catalyst.mitigate_with_zne(circuit, scale_factors=[], folding="random")()
 
     with pytest.raises(NotImplementedError):
         catalyst.qjit(mitigated_qnode)
 
 
 @pytest.mark.parametrize("params", [0.1, 0.2, 0.3, 0.4, 0.5])
 @pytest.mark.parametrize("extrapolation", [quadratic_extrapolation, exponential_extrapolate])
-def test_zne_usage_patterns(params, extrapolation):
+@pytest.mark.parametrize("folding", ["global", "all"])
+def test_zne_usage_patterns(params, extrapolation, folding):
     """Test usage patterns of catalyst.zne."""
     skip_if_exponential_extrapolation_unstable(params, extrapolation)
 
@@ -264,13 +274,13 @@ def fn(x):
     @catalyst.qjit
     def mitigated_qnode_fn_as_argument(args):
         return catalyst.mitigate_with_zne(
-            fn, scale_factors=jax.numpy.array([1, 2, 3]), extrapolate=extrapolation
+            fn, scale_factors=jax.numpy.array([1, 2, 3]), extrapolate=extrapolation, folding=folding
         )(args)
 
     @catalyst.qjit
     def mitigated_qnode_partial(args):
         return catalyst.mitigate_with_zne(
-            scale_factors=jax.numpy.array([1, 2, 3]), extrapolate=extrapolation
+            scale_factors=jax.numpy.array([1, 2, 3]), extrapolate=extrapolation, folding=folding
         )(fn)(args)
 
     assert np.allclose(mitigated_qnode_fn_as_argument(params), fn(params))

mlir/include/Mitigation/IR/MitigationOps.td

@@ -27,8 +27,8 @@ def Folding : I32EnumAttr<"Folding",
     "Folding types",
     [
         I32EnumAttrCase<"global", 1>,
-        I32EnumAttrCase<"random", 2>,
-        I32EnumAttrCase<"all",    3>,
+        I32EnumAttrCase<"all",    2>,
+        I32EnumAttrCase<"random", 3>,
     ]> {
     let cppNamespace = "catalyst::mitigation";
     let genSpecializedAttr = 0;