Implement ZNE local-random unitary folding

doc/releases/changelog-dev.md

@@ -2,6 +2,11 @@
 
 <h3>New features since last release</h3>
 
+* The `local-random` unitary folding option for :func:`~.mitigate_with_zne` is now implemented,
+  reproducing Mitiq's ``fold_gates_at_random``: every gate is folded ``floor((scale_factor-1)/2)``
+  times, then a random subset is folded once more (without replacement) to reach ``scale_factor * n``
+  gates. Non-integer scale factors are now also accepted for `local-random`.
+  [(#2956)](https://github.com/PennyLaneAI/catalyst/pull/2956)
 
 <h3>Improvements 🛠</h3>
 
@@ -361,6 +366,7 @@ Yushao Chen,
 Lillian Frederiksen,
 Sengthai Heng,
 David Ittah,
+Jacob Kitchen,
 Christina Lee,
 Mehrdad Malekmohammadi,
 River McCubbin,

frontend/catalyst/api_extensions/error_mitigation.py

@@ -48,6 +48,32 @@ def _check_is_odd_positive(numbers_list):
             raise ValueError(msg)
 
 
+def _check_is_real_ge_one(numbers_list):
+    for n in numbers_list:
+        if not isinstance(n, (int, float)):
+            msg = f"Found non-numeric {n} in scale_factors {numbers_list}.\n"
+            msg += "local-random folding requires real scale_factors >= 1"
+            raise TypeError(msg)
+        if n < 1:
+            msg = f"Found {n} < 1 in scale_factors {numbers_list}.\n"
+            msg += "local-random folding requires real scale_factors >= 1"
+            raise ValueError(msg)
+
+
+def _check_scale_factors(scale_factors, folding):
+    """Validate scale_factors against the folding method.
+
+    ``local-random`` accepts any real scale factor >= 1 (the fractional part is
+    realized at run time as a probabilistic extra fold). All other folding
+    methods scale the circuit by an exact integer factor and therefore still
+    require odd positive integers.
+    """
+    if folding == "local-random":
+        _check_is_real_ge_one(scale_factors)
+    else:
+        _check_is_odd_positive(scale_factors)
+
+
 ## API ##
 def mitigate_with_zne(
     fn=None, *, scale_factors, extrapolate=None, extrapolate_kwargs=None, folding="global"
@@ -65,7 +91,9 @@ def mitigate_with_zne(
 
     Args:
         fn (qp.QNode): the circuit to be mitigated.
-        scale_factors (list[int]): the range of noise scale factors used.
+        scale_factors (list[int] | list[float]): the range of noise scale factors used. Must be
+            odd positive integers for ``global``/``local-all`` folding. ``local-random`` folding
+            additionally accepts non-integer real scale factors >= 1.
         extrapolate (Callable): A qjit-compatible function taking two sequences as arguments (scale
             factors, and results), and returning a float by performing a fitting procedure.
             By default, perfect polynomial fitting :func:`~.polynomial_extrapolate` will be used,
@@ -75,6 +103,14 @@ def mitigate_with_zne(
         folding (str): Unitary folding technique to be used to scale the circuit. Possible values:
             - global: the global unitary of the input circuit is folded
             - local-all: per-gate folding sequences replace original gates in-place in the circuit
+            - local-random: reproduces Mitiq's ``fold_gates_at_random``. Every gate is folded
+              ``base = floor((scale_factor-1)/2)`` times, and then exactly
+              ``k = round(((scale_factor-1)/2 - base) * n)`` of the ``n`` gates are folded once
+              more, chosen uniformly at random *without replacement* (using the runtime PRNG,
+              reproducible when ``qjit(seed=...)`` is set). Odd-integer scale factors give
+              ``k == 0``, reducing to ``local-all`` and scaling the gate count exactly by
+              ``scale_factor``; non-integer scale factors (also accepted here) scale it
+              approximately by ``scale_factor``.
 
     Returns:
         Callable: A callable object that computes the mitigated of the wrapped :class:`~.QNode`
@@ -155,7 +191,7 @@ def workflow(weights, s):
     elif extrapolate_kwargs is not None:
         extrapolate = functools.partial(extrapolate, **extrapolate_kwargs)
 
-    _check_is_odd_positive(scale_factors)
+    _check_scale_factors(scale_factors, folding)
 
     return ZNECallable(fn, scale_factors, extrapolate, folding)
 
@@ -205,9 +241,6 @@ def __call__(self, *args, **kwargs):
             folding = Folding(self.folding)
         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.RANDOM:
-            raise NotImplementedError(f"Folding type {folding.value} is being developed")
 
         # Certain callables, like QNodes, may introduce additional wrappers during tracing.
         # Make sure to grab the top-level callable object in the traced function.
@@ -221,7 +254,15 @@ def __call__(self, *args, **kwargs):
             callable_fn
         ), "expected callable set as param on the first operation in zne target"
 
-        fold_numbers = (jnp.asarray(self.scale_factors, dtype=int) - 1) // 2
+        # Number of per-gate folds is (scale_factor - 1) / 2. For integer folding
+        # methods this is an exact integer count. For ``local-random`` we keep it as
+        # a float so the fractional remainder survives to the runtime, where it
+        # becomes the probability of an extra fold per gate (matching the
+        # ``scale_factor * n`` gate count Mitiq targets for fractional factors).
+        if self.folding == "local-random":
+            fold_numbers = (jnp.asarray(self.scale_factors, dtype=float) - 1) / 2
+        else:
+            fold_numbers = (jnp.asarray(self.scale_factors, dtype=int) - 1) // 2
         fold_results = zne_p.bind(
             *args_data, fold_numbers, folding=folding, jaxpr=jaxpr, fn=callable_fn
         )

frontend/test/pytest/test_mitigation.py

@@ -264,24 +264,90 @@ def mitigated_qnode():
         catalyst.qjit(mitigated_qnode)
 
 
-def test_folding_type_not_implemented():
-    """Test value of folding argument supported but not yet developed"""
+def test_local_random_folding_runs():
+    r"""Test that the local-random folding type compiles and runs.
+
+    On a noiseless simulator every folding pair $G G^\dagger$ is the identity, so the
+    mitigated result must match the unfolded circuit regardless of which gates
+    the runtime randomly selects for folding.
+    """
     dev = qp.device("lightning.qubit", wires=2)
 
     @qp.qnode(device=dev)
     def circuit():
-        return 0.0
+        qp.Hadamard(wires=0)
+        qp.RZ(0.3, wires=0)
+        qp.CNOT(wires=[1, 0])
+        qp.Hadamard(wires=1)
+        return qp.expval(qp.PauliY(wires=0))
 
+    @catalyst.qjit(seed=42)
     def mitigated_qnode():
-        return catalyst.mitigate_with_zne(circuit, scale_factors=[], folding="local-random")()
+        return catalyst.mitigate_with_zne(
+            circuit, scale_factors=[1, 3, 5], folding="local-random"
+        )()
 
-    with pytest.raises(NotImplementedError):
-        catalyst.qjit(mitigated_qnode)
+    assert np.allclose(mitigated_qnode(), circuit())
+
+
+def test_local_random_fractional_scale_factors_run():
+    r"""local-random folding accepts non-integer scale factors >= 1.
+
+    For a non-integer scale factor a random subset of gates is folded one extra time at
+    run time. On a noiseless simulator every folding pair $G G^\dagger$ is the identity, so
+    the mitigated result must still match the unfolded circuit regardless of the subset.
+    """
+    dev = qp.device("lightning.qubit", wires=2)
+
+    @qp.qnode(device=dev)
+    def circuit():
+        qp.Hadamard(wires=0)
+        qp.RZ(0.3, wires=0)
+        qp.CNOT(wires=[1, 0])
+        qp.Hadamard(wires=1)
+        return qp.expval(qp.PauliY(wires=0))
+
+    @catalyst.qjit(seed=42)
+    def mitigated_qnode():
+        return catalyst.mitigate_with_zne(
+            circuit, scale_factors=[1.0, 2.0, 3.0], folding="local-random"
+        )()
+
+    assert np.allclose(mitigated_qnode(), circuit())
+
+
+def test_local_random_rejects_scale_factor_below_one():
+    """local-random folding requires real scale factors >= 1."""
+    dev = qp.device("lightning.qubit", wires=1)
+
+    @qp.qnode(device=dev)
+    def circuit():
+        qp.Hadamard(wires=0)
+        return qp.expval(qp.PauliZ(wires=0))
+
+    with pytest.raises(ValueError, match=".*local-random folding requires real scale_factors >= 1"):
+        catalyst.mitigate_with_zne(circuit, scale_factors=[0.5], folding="local-random")
+
+
+@pytest.mark.parametrize("folding", ["global", "local-all"])
+def test_fractional_scale_factors_rejected_for_integer_folding(folding):
+    """Non-integer scale factors are only supported with local-random folding."""
+    dev = qp.device("lightning.qubit", wires=1)
+
+    @qp.qnode(device=dev)
+    def circuit():
+        qp.Hadamard(wires=0)
+        return qp.expval(qp.PauliZ(wires=0))
+
+    with pytest.raises(
+        TypeError, match=".*Only odd positive integers are allowed in scale_factors"
+    ):
+        catalyst.mitigate_with_zne(circuit, scale_factors=[1.0, 2.0], folding=folding)
 
 
 @pytest.mark.parametrize("params", [0.1, 0.2, 0.3, 0.4, 0.5])
 @pytest.mark.parametrize("extrapolation", [quadratic_extrapolation, exponential_extrapolate])
-@pytest.mark.parametrize("folding", ["global", "local-all"])
+@pytest.mark.parametrize("folding", ["global", "local-all", "local-random"])
 def test_zne_usage_patterns(params, extrapolation, folding):
     """Test usage patterns of catalyst.zne."""
     skip_if_exponential_extrapolation_unstable(params, extrapolation, threshold=0.2)

mlir/include/Mitigation/IR/MitigationOps.td

@@ -48,7 +48,12 @@ def ZneOp : Mitigation_Op<"zne", [DeclareOpInterfaceMethods<CallOpInterface>,
         SymbolRefAttr:$callee,
         Variadic<AnyType>:$args,
         FoldingAttr:$folding,
-        RankedTensorOf<[AnySignlessIntegerOrIndex]>:$numFolds,
+        // Integer entries hold an exact per-circuit fold count `(scale_factor - 1) / 2`
+        // (used by `global`/`all`). Floating-point entries carry the same quantity but
+        // keep its fractional part so `random` folding can support non-integer
+        // scale factors (base folds + one extra fold sampled with probability = frac).
+        AnyTypeOf<[RankedTensorOf<[AnySignlessIntegerOrIndex]>,
+                   RankedTensorOf<[AnyFloat]>]>:$numFolds,
         OptionalAttr<DictArrayAttr>:$arg_attrs,
         OptionalAttr<DictArrayAttr>:$res_attrs
     );