Repository: PennyLaneAI/pennylane
Pull request #7447 (2025-05-16 21:55:14)

[FTQC] Add offline byproduct correction support 
### Before submitting

Please complete the following checklist when submitting a PR:

- [x] All new features must include a unit test.
      If you've fixed a bug or added code that should be tested, add a test to the
      test directory!

- [x] All new functions and code must be clearly commented and documented.
      If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [x] Ensure that the test suite passes, by running `make test`.

- [x] Add a new entry to the `doc/releases/changelog-dev.md` file, summarizing the
      change, and including a link back to the PR.

- [x] The PennyLane source code conforms to
      [PEP8 standards](https://www.python.org/dev/peps/pep-0008/).
      We check all of our code against [Pylint](https://www.pylint.org/).
      To lint modified files, simply `pip install pylint`, and then
      run `pylint pennylane/path/to/file.py`.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

------------------------------------------------------------------------------------------------------------

**Context:**

[sc-91491]

**Background**

The `offline_byprod_correction` method allows us to move all `byproduct ops` (`X`, `Z` or `XZ`) as well as `Pauli` operations (`I`, `X`, `Y`, `Z`) of MBQC (Measurement Based Quantum Computing) to the end of a `Clifford` circuit and the measurement result can be corrected per shot based on the `Paulis` recorded in the function, i.e. the phase of `Z-basis` measurement result would be flipped if there is a `X` recorded at the target wire. 

This function reduces the number of quantum gate operations for the MBQC formalism by utilization of classical hardwares. This is achieved based on the fact that a Cliiford gate ($C$) can conjugate a Pauli to a new Pauli, i.e., $P_{new}C = CP_{old}$. 

Note that `RZ` or `RotXZX` operators at the beginning of a circuit for each wire is supported, given that no Paulis are commuted with those operators. This would be important for the `StatePrep` operations.

**Implementations**

1. The `offline_byprod_correction` method accepts `2` arguments:  
- a `tape` contains only $qml.I$, $qml.X$, $qml.Y$, $qml.Z$, $qml.S$, $qml.H$ and $qml.CNOT$; 
- mid-circuit measurement results per shot.

2. The `offline_byprod_correction` method initializes three arrays: 
- `X` correction; 
- `Z` correction; 
- `phase` information for correction of each wire.

3. Then the `tape` and `mid_meas` mid-circuit measurement results are parsed to build stack of `byprod` ops and the tape ops.

4. Then the commutate rules are applied to the byprod and the Paulis or Clifford gates in the circuit.

5. Finally, the phase factor used to correct measurement results is returned based on the recorded Paulis and measurement obs of the target wire. 

Note that a few integration tests, which are subject to be updated once the transform pipeline is done, are added for circuits with Clifford and Paulis. 

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7447?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `94.52055%` with `4 lines` in your changes missing coverage. Please review.
> Project coverage is 99.69%. Comparing base [(`a0e7e38`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/a0e7e38223c9726c9841df9ba03d383821fc94dd?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`35e0be5`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/35e0be5ae62df6428f00b9136092465ed2ba2653?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 18 commits behind head on ftqc_add_clifford_conjugate.

| [Files with missing lines](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7447?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/ftqc/pauli\_tracker.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7447?src=pr&el=tree&filepath=pennylane%2Fftqc%2Fpauli_tracker.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2Z0cWMvcGF1bGlfdHJhY2tlci5weQ==) | 94.44% | [4 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7447?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

<details><summary>Additional details and impacted files</summary>


```diff
@@                       Coverage Diff                       @@
##           ftqc_add_clifford_conjugate    #7447      +/-   ##
===============================================================
- Coverage                        99.69%   99.69%   -0.01%     
===============================================================
  Files                              538      538              
  Lines                            51267    51340      +73     
===============================================================
+ Hits                             51112    51181      +69     
- Misses                             155      159       +4     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7447?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
multiphaseCFD: > we should expect the output of our execution at this point in the pipeline is just raw, completely unprocessed samples (0s and 1s), so we should be flipping 0s and 1s, not correcting eigenvalues. We can assume the sample may have multiple wires, but that there is only a single sample measurement.

Update the correction function

> this seems to be mostly missing unit-tests, and the integration tests seem to depend on adding a rough implementation of the upcoming transform to the tests - better just to implement the transform and add integration tests for it, I think. So we can stick to just unit tests in this PR.

Removed the raw integration tests.

Patch:
diff --git a/doc/code/qml_ftqc.rst b/doc/code/qml_ftqc.rst
index 4fb293dd1..64d0ad47f 100644
--- a/doc/code/qml_ftqc.rst
+++ b/doc/code/qml_ftqc.rst
@@ -27,4 +27,5 @@ This module contains functions for tracking, commuting Pauli operations in a Cli
     ~pauli_to_xz
     ~xz_to_pauli
     ~pauli_prod
-    ~commute_clifford_op
\ No newline at end of file
+    ~commute_clifford_op
+    ~get_byproduct_corrections
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e729b3970..a3829454d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -389,6 +389,9 @@
   accessible via the `commute_clifford_op` function.
   [(#7444)](https://github.com/PennyLaneAI/pennylane/pull/7444)
 
+* Add offline byproduct correction support to the `ftqc` module.
+  [(#7447)](https://github.com/PennyLaneAI/pennylane/pull/7447)
+
 * The `ftqc` module `measure_arbitrary_basis`, `measure_x` and `measure_y` functions
   can now be captured when program capture is enabled.
   [(#7219)](https://github.com/PennyLaneAI/pennylane/pull/7219)
diff --git a/pennylane/ftqc/__init__.py b/pennylane/ftqc/__init__.py
index 11bfe0be7..eda44913c 100644
--- a/pennylane/ftqc/__init__.py
+++ b/pennylane/ftqc/__init__.py
@@ -42,6 +42,7 @@ from .lattice import Lattice, generate_lattice
 from .graph_state_preparation import GraphStatePrep, make_graph_state
 from .qubit_graph import QubitGraph
 from .utils import QubitMgr
+from .pauli_tracker import get_byproduct_corrections
 
 __all__ = [
     "GraphStatePrep",
@@ -62,4 +63,5 @@ __all__ = [
     "convert_to_mbqc_gateset",
     "generate_lattice",
     "make_graph_state",
+    "get_byproduct_corrections",
 ]
diff --git a/pennylane/ftqc/pauli_tracker.py b/pennylane/ftqc/pauli_tracker.py
index 7676fd16d..ce4f92c9b 100644
--- a/pennylane/ftqc/pauli_tracker.py
+++ b/pennylane/ftqc/pauli_tracker.py
@@ -15,12 +15,18 @@
 r"""
 This module contains Pauli Tracking functions.
 """
-
 import itertools
 from typing import List, Tuple
 
-from pennylane import CNOT, H, I, S, X, Y, Z
+import numpy as np
+
+from pennylane import math
 from pennylane.operation import Operator
+from pennylane.ops import CNOT, RZ, H, I, S, X, Y, Z
+from pennylane.tape import QuantumScript
+
+from .decomposition import _cnot_xz_corrections, _single_xz_corrections
+from .operations import RotXZX
 
 _OPS_TO_XZ = {
     I: (0, 0),
@@ -38,6 +44,10 @@ _XZ_TO_OPS = {
 
 _PAULIS = (X, Y, Z, I)
 
+_CLIFFORD_GATES_SUPPORTED = (H, S, CNOT)
+
+_NON_CLIFFORD_GATES_SUPPORTED = (RZ, RotXZX)
+
 
 def pauli_to_xz(op: Operator) -> Tuple[int, int]:
     r"""
@@ -235,3 +245,250 @@ def commute_clifford_op(clifford_op: Operator, xz: List[Tuple[int, int]]) -> Lis
         return _commute_cnot(_xc, _zc, _xt, _zt)
 
     raise NotImplementedError("Only qml.H, qml.S and qml.CNOT are supported.")
+
+
+def _parse_mid_measurements(tape: QuantumScript, mid_meas: List):
+    r"""Parse a series of mid-measurement results of a quantum tape where
+    all gates are from the set {X, Y, Z, I, H, S, CNOT, RZ, RotXZX}, and
+    where only the first gate on each wire is permitted to be a non-Clifford
+    gates (RZ, RotXZX). Both non-Clifford and Clifford gates mentioned above
+    are measured in the way defined in Raussendorf et al. <https://arxiv.org/abs/quant-ph/0301052>`__.
+
+    Args:
+        tape (QuantumScript): The quantum tape in the standard circuit mode (Gates are not transformed into the MBQC formalism).
+        mid_meas (list): Mid-measurements results.
+
+    Returns:
+        A list of `byproduct` ops in a reversed manner.
+    """
+    ops = tape.operations
+
+    by_ops = []
+
+    mid_meas_offset = 0
+    _wires_used = [0] * (max(tape.wires) + 1)
+    for op in ops:
+        for wire in op.wires:
+            _wires_used[wire] += 1
+
+        gate_offset = 4 if op.num_wires == 1 else 13
+        ms = mid_meas[mid_meas_offset : mid_meas_offset + gate_offset]
+        by_op = []
+        if isinstance(op, (S, H)):
+            by_op = [_single_xz_corrections(op, *ms)]
+        elif isinstance(op, CNOT):
+            by_op = _cnot_xz_corrections(ms)
+        elif isinstance(op, (RZ, RotXZX)):
+            if _wires_used[op.wires[0]] > 1:
+                raise NotImplementedError(
+                    "The current implementation requires that only a single non-Clifford gate comes before Clifford and Pauli gates for each wire."
+                )
+            by_op = [_single_xz_corrections(op, *ms)]
+        elif isinstance(op, _PAULIS):
+            continue
+        else:
+            raise NotImplementedError(f"{op.name} is not supported.")
+
+        by_ops.append(by_op)
+
+        mid_meas_offset += gate_offset
+
+    # To use list in a stack manner
+    by_ops.reverse()
+    return by_ops
+
+
+def _get_xz_record(tape: QuantumScript, by_ops: List[Tuple[int, int]]):
+    """Commutate/merge the Pauli/byproduct ops of a Clifford circuit.
+
+    Args:
+        tape (QuantumScript): A quantum tape.
+        by_ops (list): List of byproduct operators for Clifford gates
+
+    Return:
+        The final recorded x and z for each wire.
+    """
+
+    num_wires = max(tape.wires) + 1
+
+    x_record = math.zeros(num_wires, dtype=np.uint8)
+    z_record = math.zeros(num_wires, dtype=np.uint8)
+
+    for op in tape.operations:
+        wires = list(op.wires)
+
+        # Get the recorded xz
+        xz = [(x_record[wire], z_record[wire]) for wire in wires]
+
+        # Updated xz
+        new_xz = []
+        if isinstance(op, _NON_CLIFFORD_GATES_SUPPORTED):
+            # Branch for non-Clifford gates
+            new_xz.append(by_ops.pop()[0])
+        elif isinstance(op, _CLIFFORD_GATES_SUPPORTED):
+            # Branch for Clifford gates
+            # Step 1: Commutate the recorded xz with the Clifford gate to a new xz.
+            xz_commutated = commute_clifford_op(op, xz)
+
+            # Step 2: Merge the new xz with the byproduct by_op
+            by_op = by_ops.pop()
+            for _by_op, _xz_comm in zip(by_op, xz_commutated):
+                new_xz.append(math.bitwise_xor(_by_op, _xz_comm))
+        else:  # branch for Paulis
+            # Commutate step is skipped.
+            # Get the new xz by merging the recorded xz with the Pauli ops directly.
+            new_xz.append(math.bitwise_xor(pauli_to_xz(op), xz[0]))
+        # Assign the updated the xz to the x, z record
+        for idx, wire in enumerate(wires):
+            x_record[wire], z_record[wire] = new_xz[idx]
+
+    return x_record, z_record
+
+
+def _correct_samples(tape: QuantumScript, x_record: math.array, measurement_vals: List):
+    """Correct sample measurements in a tape. The samples are corrected based on the `samples`
+    at `wires` with the corresponding recorded x.
+
+        Args:
+            tape (tape: qml.tape.QuantumScript): A quantum tape.
+            measurement_vals (List) : A list of measurement values.
+            x_record (math.array): The array of recorded x for each wire.
+            measurement_vals (List): Measurement values.
+
+        Return:
+            A list of corrected measurement values.
+    """
+    measured_wires = tape.measurements[0].wires
+    x_at_target = np.array([x_record[w] for w in measured_wires])
+    correct_meas = np.bitwise_xor(np.array(measurement_vals), x_at_target)
+
+    return correct_meas
+
+
+def get_byproduct_corrections(tape: QuantumScript, mid_meas: List, measurement_vals: List):
+    r"""Correct sample results offline based on the executed quantum script and the mid-circuit measurement results for each shot.
+    The mid measurement results are first parsed with the quantum script to get the byproduct operations for each Clifford
+    and non-Clifford gates. Note that byproduct operations are stored as a list and accessed in a stack manner. The calculation iteratively
+    pops out the first operation in the tape and applies commutation rules for the first byproduct ops in the byproduct stack and
+    then the results are commutated to the byproduct of the current operations in the tape if it is a Clifford gate. The calculation
+    starts from applying commutate rules for :class:`qml.I` gate or :math:`encode\_xz(x,z)=(0,0)` to the first gate in the tape. The
+    measurement corrections are returned based on the observable operators and the xz recorded.
+
+    Args:
+        tape (tape: qml.tape.QuantumScript): A Clifford quantum tape with :class:`~pennylane.X`, :class:`~pennylane.Y`, :class:`~pennylane.Z`,
+            :class:`~pennylane.I`, :class:`~pennylane.H`, :class:`~pennylane.S`, :class:`~pennylane.CNOT` and non-Clifford gates (:class:`~pennylane.RZ`
+            and :class:`~pennylane.ftqc.RotXZX`) at the beginning of circuit in the standard circuit formalism. Note that one non-Clifford gate per wire
+            at most is supported.
+        mid_meas (list): MidMeasurement results per shot.
+        measurement_vals (list): Raw measurement results.
+
+    Return:
+        A list of corrected measurement results.
+
+
+    **Note**
+    This work is to be integrated into the MBQC transform pipeline.
+
+    **Example:**
+
+        .. code-block:: python3
+
+            import pennylane as qml
+
+            from pennylane.ftqc import diagonalize_mcms, generate_lattice, measure_x, measure_y
+            from pennylane.ftqc import GraphStatePrep
+
+            from pennylane.ftqc.pauli_tracker import get_byproduct_corrections
+            import numpy as np
+
+
+            def generate_random_state(n):
+                seed_value = 42  # You can use any integer as the seed
+                np.random.seed(seed_value)
+                input_state = np.random.random(2**n) + 1j * np.random.random(2**n)
+                return input_state / np.linalg.norm(input_state)
+
+
+            def generate_rot_gate_graph():
+                lattice = generate_lattice([4], "chain")
+                return lattice.graph
+
+
+            num_shots = 1000
+            dev = qml.device("lightning.qubit", shots=num_shots)
+
+            @diagonalize_mcms
+            @qml.qnode(dev, mcm_method="one-shot")
+            def circ(start_state):
+                qml.StatePrep(start_state, wires=[0])
+                GraphStatePrep(generate_rot_gate_graph(), wires=[1, 2, 3, 4])
+                qml.CZ(wires=[0, 1])
+                m0 = measure_x(0, reset=True)
+                m1 = measure_y(1, reset=True)
+                m2 = measure_y(2, reset=True)
+                m3 = measure_y(3, reset=True)
+
+                GraphStatePrep(generate_rot_gate_graph(), wires=[3, 2, 1, 0])
+                qml.CZ(wires=[3, 4])
+                m4 = measure_x(4, reset=True)
+                m5 = measure_y(3, reset=True)
+                m6 = measure_y(2, reset=True)
+                m7 = measure_y(1, reset=True)
+
+                GraphStatePrep(generate_rot_gate_graph(), wires=[1, 2, 3, 4])
+                qml.CZ(wires=[0, 1])
+                m8 = measure_x(0, reset=True)
+                m9 = measure_y(1, reset=True)
+                m10 = measure_y(2, reset=True)
+                m11 = measure_y(3, reset=True)
+
+                return (
+                    qml.sample(wires=[4]),
+                    qml.sample(m0),
+                    qml.sample(m1),
+                    qml.sample(m2),
+                    qml.sample(m3),
+                    qml.sample(m4),
+                    qml.sample(m5),
+                    qml.sample(m6),
+                    qml.sample(m7),
+                    qml.sample(m8), qml.sample(m9), qml.sample(m10), qml.sample(m11)
+                )
+
+            init_state = generate_random_state(1)
+            res = circ(init_state)
+
+            ops = [qml.H(wires=[0]), qml.H(wires=[0]), qml.H(wires=[0])]
+            measurements = [qml.sample(qml.Z(0))]
+
+            meas_res = res[0]
+            mid_meas_res = res[1:]
+            corrected_meas_res = []
+
+            script = qml.tape.QuantumScript(ops, measurements, shots=num_shots)
+
+            for i in range(num_shots):
+                mid_meas = [row[i] for row in mid_meas_res]
+                corrected_meas_res.extend(get_byproduct_corrections(script, mid_meas, [meas_res[i]]))
+
+            res_corrected = 1 - 2*np.sum(corrected_meas_res) / num_shots
+
+            dev_ref = qml.device("default.qubit")
+
+            @diagonalize_mcms
+            @qml.qnode(dev)
+            def circ_ref(start_state):
+                qml.StatePrep(start_state, wires=[0])
+                qml.H(0)
+                qml.H(0)
+                qml.H(0)
+                return qml.expval(qml.Z(0))
+
+            np.allclose(res_corrected, circ_ref(init_state))
+
+    """
+    by_ops = _parse_mid_measurements(tape, mid_meas)
+
+    x_record, _ = _get_xz_record(tape, by_ops)
+
+    return _correct_samples(tape, x_record, measurement_vals)


Test patch:
diff --git a/tests/ftqc/test_pauli_tracker.py b/tests/ftqc/test_pauli_tracker.py
index 261fd376d..80867640a 100644
--- a/tests/ftqc/test_pauli_tracker.py
+++ b/tests/ftqc/test_pauli_tracker.py
@@ -19,7 +19,14 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.ftqc.pauli_tracker import commute_clifford_op, pauli_prod, pauli_to_xz, xz_to_pauli
+from pennylane.ftqc import RotXZX, get_byproduct_corrections
+from pennylane.ftqc.pauli_tracker import (
+    _parse_mid_measurements,
+    commute_clifford_op,
+    pauli_prod,
+    pauli_to_xz,
+    xz_to_pauli,
+)
 
 _PAULIS = (qml.I, qml.X, qml.Y, qml.Z)
 
@@ -33,6 +40,12 @@ def generate_pauli_list(wire: int, num_ops: int):
     return pauli_list
 
 
+def generate_random_state(n, seed=0):
+    rng = np.random.default_rng(seed=seed)
+    input_state = rng.random(2**n) + 1j * rng.random(2**n)
+    return input_state / np.linalg.norm(input_state)
+
+
 class TestPauliTracker:
     """Test for the pauli tracker related functions."""
 
@@ -186,3 +199,120 @@ class TestPauliTracker:
             match="Please ensure xz are either 0 or 1.",
         ):
             _ = commute_clifford_op(clifford_op, xz)
+
+
+class TestOfflineCorrection:
+    """Tests for byproduct operation offline corrections."""
+
+    @pytest.mark.parametrize(
+        "ops, mid_measures, expected",
+        [
+            (
+                [qml.CNOT(wires=[0, 1]), qml.X(0), qml.X(1)],
+                [1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 1, 0],
+                [[(0, 1), (0, 1)]],
+            ),
+            ([qml.S(0), qml.Y(0)], [1, 0, 0, 1], [[(1, 0)]]),
+            ([qml.H(0), qml.Z(0)], [1, 1, 0, 0], [[(1, 1)]]),
+            ([qml.RZ(0.1, wires=[0]), qml.X(0)], [1, 0, 1, 1], [[(1, 0)]]),
+        ],
+    )
+    def test_parse_mid_measurements(self, ops, mid_measures, expected):
+        tape = qml.tape.QuantumScript(ops=ops)
+        by_ops = _parse_mid_measurements(tape, mid_measures)
+        assert by_ops == expected
+
+    @pytest.mark.parametrize("num_shots", [250])
+    @pytest.mark.parametrize(
+        "ops",
+        [
+            [qml.RZ(0.1, wires=[0]), qml.RZ(0.3, wires=[0])],
+            [qml.RZ(0.1, wires=[0]), RotXZX(0.1, 0.2, 0.3, wires=[0])],
+            [qml.MultiControlledX(wires=[0, 1, 2, 3], control_values=[0, 1, 0])],
+        ],
+    )
+    def test_unsupported_tape_offline_correction(self, num_shots, ops):
+        measurements = [qml.sample(wires=[0])]
+
+        tape = qml.tape.QuantumScript(ops=ops, measurements=measurements, shots=num_shots)
+
+        mid_meas = [random.choice([0, 1]) for _ in range(8)]
+
+        with pytest.raises(NotImplementedError):
+            _ = get_byproduct_corrections(tape, mid_meas, [1])
+
+    @pytest.mark.parametrize(
+        "ops, mid_meas, sample_wires, measures_expected",
+        [
+            (
+                [qml.RZ(0.1, wires=[0]), qml.RZ(0.3, wires=[1]), qml.Z(0)],
+                [0, 1, 1, 0, 1, 1, 0, 0],
+                [0],
+                [[1], [0]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), qml.RZ(0.3, wires=[2]), qml.Z(0)],
+                [0, 1, 1, 0, 1, 1, 0, 0],
+                [0],
+                [[1], [0]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), qml.S(wires=[0])],
+                [0, 0, 1, 1, 1, 1, 0, 0],
+                [0],
+                [[1], [1]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), RotXZX(0.1, -0.1, 0.3, wires=[1])],
+                [0, 1, 0, 1, 1, 0, 0, 1],
+                [0],
+                [[0], [0]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), RotXZX(0.1, -0.1, 0.3, wires=[1])],
+                [0, 0, 1, 0, 1, 0, 0, 1],
+                [0],
+                [[1], [1]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), qml.RZ(0.3, wires=[1])],
+                [0, 1, 1, 0, 1, 1, 0, 0],
+                [1],
+                [[0], [1]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), RotXZX(0.1, 0.2, 0.3, wires=[1])],
+                [0, 1, 1, 0, 1, 0, 0, 0],
+                [1],
+                [[0], [0]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), RotXZX(0.1, -0.1, 0.3, wires=[1])],
+                [0, 1, 0, 1, 1, 0, 0, 1],
+                [1],
+                [[0], [1]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), RotXZX(0.1, -0.1, 0.3, wires=[1])],
+                [0, 1, 0, 1, 1, 1, 0, 1],
+                [1],
+                [[1], [1]],
+            ),
+            (
+                [qml.RZ(-0.1, wires=[0]), qml.RZ(0.1, wires=[1]), qml.X(wires=[1])],
+                [0, 1, 0, 1, 1, 1, 0, 1],
+                [0, 1],
+                [[1, 0], [1, 1]],
+            ),
+        ],
+    )
+    def test_non_clifford_offline_correction(self, ops, mid_meas, sample_wires, measures_expected):
+        measurements = [qml.sample(wires=sample_wires)]
+
+        tape = qml.tape.QuantumScript(ops=ops, measurements=measurements)
+
+        measures, expected = measures_expected
+
+        corrected_meas = get_byproduct_corrections(tape, mid_meas, measures)
+
+        assert np.allclose(corrected_meas, expected)