Test audit/story 2 measurements

frontend/test/pytest/conftest.py

@@ -107,7 +107,11 @@ def circuit():
         if "old_frontend" in request.keywords:
             pytest.skip("Test is specific to the old frontend and should not run with capture.")
         if "capture_todo" in request.keywords:
-            pytest.xfail("Not expected to work yet with program capture.")
+            request.node.add_marker(
+                pytest.mark.xfail(
+                    reason="Not expected to work yet with program capture.", strict=True
+                )
+            )
     return request.param
 
 

frontend/test/pytest/test_measurement_dynamic_shapes.py

@@ -88,8 +88,10 @@ def circuit():
     assert out.count("compiling...") == 1
 
 
+# capture gap: capture=True fails in measurement lowering/interpreter pathway for this scenario.
+# fix direction: close capture measurement gap in from_plxpr/qfunc_interpreter and normalize behavior with legacy execution.
 @pytest.mark.parametrize("readout", [qml.expval, qml.var])
-def test_dynamic_wires_scalar_readouts(readout, backend, capfd):
+def test_dynamic_wires_scalar_readouts(readout, backend, capfd, capture_mode):
     """
     Test that a circuit with dynamic number of wires can be executed correctly.
 
@@ -103,7 +105,7 @@ def ref(num_qubits):
 
         @qml.qnode(dev)
         def circ():
-            @catalyst.for_loop(0, num_qubits, 1)
+            @qml.for_loop(0, num_qubits, 1)
             def loop_0(i):
                 qml.RY(2.2, wires=i)
 
@@ -113,16 +115,18 @@ def loop_0(i):
 
         return circ()
 
-    cat = catalyst.qjit(ref)
+    cat = catalyst.qjit(ref, capture=capture_mode)
 
     assert np.allclose(ref(10), cat(10))
     assert np.allclose(ref(4), cat(4))
     out, _ = capfd.readouterr()
     assert out.count("compiling...") == 3
 
 
+# capture gap: capture=True fails in measurement lowering/interpreter pathway for this scenario.
+# fix direction: close capture measurement gap in from_plxpr/qfunc_interpreter and normalize behavior with legacy execution.
 @pytest.mark.parametrize("readout", [qml.probs])
-def test_dynamic_wires_statebased_with_wires(readout, backend, capfd):
+def test_dynamic_wires_statebased_with_wires(readout, backend, capfd, capture_mode):
     """
     Test that a circuit with dynamic number of wires can be executed correctly
     with state based measurements with wires specified.
@@ -136,7 +140,7 @@ def ref(num_qubits):
 
         @qml.qnode(dev)
         def circ():
-            @catalyst.for_loop(0, num_qubits, 1)
+            @qml.for_loop(0, num_qubits, 1)
             def loop_0(i):
                 qml.RY(2.2, wires=i)
 
@@ -148,16 +152,18 @@ def loop_0(i):
 
         return circ()
 
-    cat = catalyst.qjit(ref)
+    cat = catalyst.qjit(ref, capture=capture_mode)
 
     assert np.allclose(ref(10), cat(10))
     assert np.allclose(ref(4), cat(4))
     out, _ = capfd.readouterr()
     assert out.count("compiling...") == 3
 
 
+# capture gap: capture=True fails in measurement lowering/interpreter pathway for this scenario.
+# fix direction: close capture measurement gap in from_plxpr/qfunc_interpreter and normalize behavior with legacy execution.
 @pytest.mark.parametrize("readout", [qml.probs, qml.state])
-def test_dynamic_wires_statebased_without_wires(readout, backend, capfd):
+def test_dynamic_wires_statebased_without_wires(readout, backend, capfd, capture_mode):
     """
     Test that a circuit with dynamic number of wires can be executed correctly
     with state based measurements without wires specified.
@@ -169,7 +175,7 @@ def ref(num_qubits):
 
         @qml.qnode(dev)
         def circ(x):
-            @catalyst.for_loop(0, num_qubits, 1)
+            @qml.for_loop(0, num_qubits, 1)
             def loop_0(i):
                 qml.RY(2.2, wires=i)
 
@@ -179,16 +185,18 @@ def loop_0(i):
 
         return circ(42)
 
-    cat = catalyst.qjit(ref)
+    cat = catalyst.qjit(ref, capture=capture_mode)
 
     assert np.allclose(ref(10), cat(10))
     assert np.allclose(ref(4), cat(4))
     out, _ = capfd.readouterr()
     assert out.count("compiling...") == 3
 
 
+# capture gap: capture=True fails in measurement lowering/interpreter pathway for this scenario.
+# fix direction: close capture measurement gap in from_plxpr/qfunc_interpreter and normalize behavior with legacy execution.
 @pytest.mark.parametrize("shots", [3, (3, 4, 5)])
-def test_dynamic_wires_sample_with_wires(shots, backend, capfd):
+def test_dynamic_wires_sample_with_wires(shots, backend, capfd, capture_mode):
     """
     Test that a circuit with dynamic number of wires can be executed correctly
     with sample measurements with wires specified.
@@ -201,7 +209,7 @@ def ref(num_qubits):
         @qml.set_shots(shots)
         @qml.qnode(dev)
         def circ():
-            @catalyst.for_loop(0, num_qubits, 1)
+            @qml.for_loop(0, num_qubits, 1)
             def loop_0(i):
                 qml.RY(0.0, wires=i)
 
@@ -211,7 +219,7 @@ def loop_0(i):
 
         return circ()
 
-    cat = catalyst.qjit(ref)
+    cat = catalyst.qjit(ref, capture=capture_mode)
     num_shots = 1 if isinstance(shots, int) else len(shots)
     for test_nqubits in (10, 4):
         expected = ref(test_nqubits)
@@ -222,8 +230,10 @@ def loop_0(i):
     assert out.count("compiling...") == 3
 
 
+# capture gap: capture=True fails in measurement lowering/interpreter pathway for this scenario.
+# fix direction: close capture measurement gap in from_plxpr/qfunc_interpreter and normalize behavior with legacy execution.
 @pytest.mark.parametrize("shots", [3, (3, 4, 5), (7,) * 3])
-def test_dynamic_wires_sample_without_wires(shots, backend, capfd):
+def test_dynamic_wires_sample_without_wires(shots, backend, capfd, capture_mode):
     """
     Test that a circuit with dynamic number of wires can be executed correctly
     with sample measurements without wires specified.
@@ -236,7 +246,7 @@ def ref(num_qubits):
         @qml.set_shots(shots)
         @qml.qnode(dev)
         def circ():
-            @catalyst.for_loop(0, num_qubits, 1)
+            @qml.for_loop(0, num_qubits, 1)
             def loop_0(i):
                 qml.RY(0.0, wires=i)
 
@@ -246,7 +256,7 @@ def loop_0(i):
 
         return circ()
 
-    cat = catalyst.qjit(ref)
+    cat = catalyst.qjit(ref, capture=capture_mode)
     num_shots = 1 if isinstance(shots, int) else len(shots)
     for test_nqubits in (10, 4):
         expected = ref(test_nqubits)
@@ -257,15 +267,17 @@ def loop_0(i):
     assert out.count("compiling...") == 3
 
 
-def test_dynamic_wires_counts_with_wires(backend, capfd):
+# capture gap: capture=True fails in measurement lowering/interpreter pathway for this scenario.
+# fix direction: close capture measurement gap in from_plxpr/qfunc_interpreter and normalize behavior with legacy execution.
+def test_dynamic_wires_counts_with_wires(backend, capfd, capture_mode):
     """
     Test that a circuit with dynamic number of wires can be executed correctly
     with counts measurements with wires specified.
 
     Note that Catalyst does not support shot vectors with counts.
     """
 
-    @catalyst.qjit
+    @catalyst.qjit(capture=capture_mode)
     def func(num_qubits):
         print("compiling...")
         dev = qml.device(backend, wires=num_qubits)
@@ -287,15 +299,17 @@ def circ():
     assert out.count("compiling...") == 1
 
 
-def test_dynamic_wires_counts_without_wires(backend, capfd):
+# capture gap: capture=True fails in measurement lowering/interpreter pathway for this scenario.
+# fix direction: close capture measurement gap in from_plxpr/qfunc_interpreter and normalize behavior with legacy execution.
+def test_dynamic_wires_counts_without_wires(backend, capfd, capture_mode):
     """
     Test that a circuit with dynamic number of wires can be executed correctly
     with counts measurements without wires specified.
 
     Note that Catalyst does not support shot vectors with counts.
     """
 
-    @catalyst.qjit
+    @catalyst.qjit(capture=capture_mode)
     def func(num_qubits):
         print("compiling...")
         dev = qml.device(backend, wires=num_qubits)
@@ -320,14 +334,16 @@ def circ():
     assert out.count("compiling...") == 1
 
 
+# capture gap: capture=True fails in measurement lowering/interpreter pathway for this scenario.
+# fix direction: close capture measurement gap in from_plxpr/qfunc_interpreter and normalize behavior with legacy execution.
 @pytest.mark.parametrize("wires", [1.1, (1.1)])
-def test_wrong_wires_argument(backend, wires):
+def test_wrong_wires_argument(backend, wires, capture_mode):
     """
     Test that a circuit with a wrongly typed and shaped dynamic wire argument
     is correctly caught.
     """
 
-    @catalyst.qjit
+    @catalyst.qjit(capture=capture_mode)
     def func(num_qubits):
         dev = qml.device(backend, wires=num_qubits)
 
@@ -343,10 +359,12 @@ def circ():
         func(wires)
 
 
-def test_dynamic_shots_and_wires(capfd):
+# capture gap: capture=True fails in measurement lowering/interpreter pathway for this scenario.
+# fix direction: close capture measurement gap in from_plxpr/qfunc_interpreter and normalize behavior with legacy execution.
+def test_dynamic_shots_and_wires(capfd, capture_mode):
     """Test that a circuit with both dynamic shots and dynamic wires works correctly with qml.sample."""
 
-    @catalyst.qjit
+    @catalyst.qjit(capture=capture_mode)
     def workflow_dynamic_shots_and_wires(num_shots, num_wires):
         print("compiling...")
         device = qml.device("lightning.qubit", wires=num_wires)
@@ -355,16 +373,16 @@ def workflow_dynamic_shots_and_wires(num_shots, num_wires):
         @qml.qnode(device)
         def circuit():
             # Apply Hadamard to all wires
-            @catalyst.for_loop(0, num_wires, 1)
+            @qml.for_loop(0, num_wires, 1)
             def apply_hadamards(i):
                 qml.Hadamard(i)
 
             apply_hadamards()
 
             # Apply some entangling gates if we have multiple wires
-            @catalyst.cond(num_wires > 1)
+            @qml.cond(num_wires > 1)
             def add_entanglement():
-                @catalyst.for_loop(0, num_wires - 1, 1)
+                @qml.for_loop(0, num_wires - 1, 1)
                 def apply_cnots(i):
                     qml.CNOT([i, i + 1])