Allow abstract Boolean measurement results as a parameter to qml.samp…

…le (#5673)

This PR allows `qml.sample` (and implicitly other measurement processes)
to accept a Boolean tracer in the observable parameter, similar to what
is already done for the `MeasurementValue` value object. The `wires`
attribute will be empty in this case.

[sc-62096]

doc/releases/changelog-dev.md

@@ -88,6 +88,11 @@
 * Sets up the framework for the development of an `assert_equal` function for testing operator comparison.
   [(#5634)](https://github.com/PennyLaneAI/pennylane/pull/5634)
 
+* `qml.sample` can now be used on Boolean values representing mid-circuit measurement results in
+  traced quantum functions. This feature is used with Catalyst to enable the pattern
+  `m = measure(0); qml.sample(m)`.
+  [(#5673)](https://github.com/PennyLaneAI/pennylane/pull/5673)
+
 * PennyLane operators can now automatically be captured as instructions in JAXPR. See the experimental
   `capture` module for more information.
   [(#5511)](https://github.com/PennyLaneAI/pennylane/pull/5511)
@@ -202,6 +207,7 @@ Ahmed Darwish,
 Isaac De Vlugt,
 Pietropaolo Frisoni,
 Emiliano Godinez,
+David Ittah,
 Soran Jahangiri,
 Korbinian Kottmann,
 Christina Lee,

pennylane/measurements/measurements.py

@@ -23,6 +23,7 @@
 from typing import Optional, Sequence, Tuple, Union
 
 import pennylane as qml
+from pennylane.math.utils import is_abstract
 from pennylane.operation import DecompositionUndefinedError, EigvalsUndefinedError, Operator
 from pennylane.pytrees import register_pytree
 from pennylane.typing import TensorLike
@@ -162,6 +163,9 @@ def __init__(
             # Cast sequence of measurement values to list
             self.mv = obs if getattr(obs, "name", None) == "MeasurementValue" else list(obs)
             self.obs = None
+        elif is_abstract(obs):  # Catalyst program with qml.sample(m, wires=i)
+            self.mv = obs
+            self.obs = None
         else:
             self.obs = obs
             self.mv = None
@@ -306,7 +310,7 @@ def wires(self):
 
         This is the union of all the Wires objects of the measurement.
         """
-        if self.mv is not None:
+        if self.mv is not None and not is_abstract(self.mv):
             if isinstance(self.mv, list):
                 return qml.wires.Wires.all_wires([m.wires for m in self.mv])
             return self.mv.wires

tests/measurements/test_sample.py

@@ -499,6 +499,24 @@ def circuit(x):
     )
 
 
+@pytest.mark.jax
+def test_sample_with_boolean_tracer():
+    """Test that qml.sample can be used with Catalyst measurement values (Boolean tracer)."""
+    import jax
+
+    def fun(b):
+        mp = qml.sample(b)
+
+        assert mp.obs is None
+        assert isinstance(mp.mv, jax.interpreters.partial_eval.DynamicJaxprTracer)
+        assert mp.mv.dtype == bool
+        assert mp.mv.shape == ()
+        assert isinstance(mp.wires, qml.wires.Wires)
+        assert mp.wires == ()
+
+    jax.make_jaxpr(fun)(True)
+
+
 @pytest.mark.jax
 @pytest.mark.parametrize(
     "obs",

tests/test_compiler.py

@@ -716,3 +716,24 @@ def f(x):
             CompileError, match="Pennylane does not support the VJP function without QJIT."
         ):
             vjp(x, dy)
+
+
+class TestCatalystSample:
+    """Test qml.sample with Catalyst."""
+
+    @pytest.mark.xfail(reason="requires simultaneous catalyst pr")
+    def test_sample_measure(self):
+        """Test that qml.sample can be used with catalyst.measure."""
+
+        dev = qml.device("lightning.qubit", wires=1, shots=1)
+
+        @qml.qjit
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RY(x, wires=0)
+            m = catalyst.measure(0)
+            qml.PauliX(0)
+            return qml.sample(m)
+
+        assert circuit(0.0) == 0
+        assert circuit(jnp.pi) == 1