dynamic_one_shot uses tapes with shot-vectors and jitting takes advantage of it

doc/releases/changelog-dev.md

@@ -6,6 +6,11 @@
 
 <h3>Improvements 🛠</h3>
 
+<h4>Mid-circuit measurements and dynamic circuits</h4>
+
+* The `dynamic_one_shot` transform can be compiled with `jax.jit`.
+  [(#5557)](https://github.com/PennyLaneAI/pennylane/pull/5557)
+
 * When using `defer_measurements` with postselecting mid-circuit measurements, operations
   that will never be active due to the postselected state are skipped in the transformed
   quantum circuit. In addition, postselected controls are skipped, as they are evaluated
@@ -77,4 +82,5 @@ This release contains contributions from (in alphabetical order):
 
 Pietropaolo Frisoni,
 Christina Lee,
-David Wierichs.
+Vincent Michaud-Rioux,
+David Wierichs.

pennylane/_qubit_device.py

@@ -275,8 +275,19 @@ def execute(self, circuit, **kwargs):
         self.check_validity(circuit.operations, circuit.observables)
 
         has_mcm = any(isinstance(op, MidMeasureMP) for op in circuit.operations)
-        if has_mcm:
-            kwargs["mid_measurements"] = {}
+        if has_mcm and "mid_measurements" not in kwargs:
+            results = []
+            aux_circ = qml.tape.QuantumScript(
+                circuit.operations,
+                circuit.measurements,
+                shots=[1],
+                trainable_params=circuit.trainable_params,
+            )
+            for _ in circuit.shots:
+                kwargs["mid_measurements"] = {}
+                self.reset()
+                results.append(self.execute(aux_circ, **kwargs))
+            return tuple(results)
         # apply all circuit operations
         self.apply(
             circuit.operations,

pennylane/devices/qubit/apply_operation.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Functions to apply an operation to a state vector."""
-# pylint: disable=unused-argument
+# pylint: disable=unused-argument, too-many-arguments
 
 from functools import singledispatch
 from string import ascii_letters as alphabet
@@ -21,7 +21,7 @@
 
 import pennylane as qml
 from pennylane import math
-from pennylane.measurements import MidMeasureMP, Shots
+from pennylane.measurements import MidMeasureMP
 from pennylane.ops import Conditional
 
 SQRT2INV = 1 / math.sqrt(2)
@@ -238,14 +238,36 @@ def apply_conditional(
         ndarray: output state
     """
     mid_measurements = execution_kwargs.get("mid_measurements", None)
-
+    rng = execution_kwargs.get("rng", None)
+    prng_key = execution_kwargs.get("prng_key", None)
+    interface = qml.math.get_deep_interface(state)
+    if interface == "jax":
+        # pylint: disable=import-outside-toplevel
+        from jax.lax import cond
+
+        return cond(
+            op.meas_val.concretize(mid_measurements),
+            lambda x: apply_operation(
+                op.then_op,
+                x,
+                is_state_batched=is_state_batched,
+                debugger=debugger,
+                mid_measurements=mid_measurements,
+                rng=rng,
+                prng_key=prng_key,
+            ),
+            lambda x: x,
+            state,
+        )
     if op.meas_val.concretize(mid_measurements):
         return apply_operation(
             op.then_op,
             state,
             is_state_batched=is_state_batched,
             debugger=debugger,
             mid_measurements=mid_measurements,
+            rng=rng,
+            prng_key=prng_key,
         )
     return state
 
@@ -273,31 +295,47 @@ def apply_mid_measure(
     mid_measurements = execution_kwargs.get("mid_measurements", None)
     rng = execution_kwargs.get("rng", None)
     prng_key = execution_kwargs.get("prng_key", None)
-
     if is_state_batched:
         raise ValueError("MidMeasureMP cannot be applied to batched states.")
-    if not np.allclose(np.linalg.norm(state), 1.0):
-        mid_measurements[op] = -1
-        return np.zeros_like(state)
     wire = op.wires
-    sample = qml.devices.qubit.sampling.measure_with_samples(
-        [qml.sample(wires=wire)], state, Shots(1), rng=rng, prng_key=prng_key
-    )
-    sample = int(sample[0])
-    mid_measurements[op] = sample
-    if op.postselect is not None and sample != op.postselect:
-        mid_measurements[op] = -1
-        return np.zeros_like(state)
     axis = wire.toarray()[0]
     slices = [slice(None)] * qml.math.ndim(state)
-    slices[axis] = int(not sample)
-    state[tuple(slices)] = 0.0
-    state_norm = np.linalg.norm(state)
-    state = state / state_norm
-    if op.reset and sample == 1:
-        state = apply_operation(
-            qml.X(wire), state, is_state_batched=is_state_batched, debugger=debugger
-        )
+    slices[axis] = 0
+    prob0 = qml.math.norm(state[tuple(slices)]) ** 2
+    interface = qml.math.get_deep_interface(state)
+    if prng_key is not None:
+        # pylint: disable=import-outside-toplevel
+        from jax.random import binomial
+
+        def binomial_fn(n, p):
+            return binomial(prng_key, n, p).astype(int)
+
+    else:
+        binomial_fn = np.random.binomial if rng is None else rng.binomial
+    sample = binomial_fn(1, 1 - prob0)
+    mid_measurements[op] = sample
+
+    # Using apply_operation(qml.QubitUnitary,...) instead of apply_operation(qml.Projector([sample], wire),...)
+    # to select the sample branch enables jax.jit and prevents it from using Python callbacks
+    matrix = qml.math.array([[(sample + 1) % 2, 0.0], [0.0, (sample) % 2]], like=interface)
+    state = apply_operation(
+        qml.QubitUnitary(matrix, wire),
+        state,
+        is_state_batched=is_state_batched,
+        debugger=debugger,
+    )
+    state = state / qml.math.norm(state)
+
+    # Using apply_operation(qml.QubitUnitary,...) instead of apply_operation(qml.X(wire), ...)
+    # to reset enables jax.jit and prevents it from using Python callbacks
+    element = op.reset and sample == 1
+    matrix = qml.math.array(
+        [[(element + 1) % 2, (element) % 2], [(element) % 2, (element + 1) % 2]], like=interface
+    ).astype(float)
+    state = apply_operation(
+        qml.QubitUnitary(matrix, wire), state, is_state_batched=is_state_batched, debugger=debugger
+    )
+
     return state
 
 

pennylane/devices/qubit/sampling.py

@@ -35,7 +35,7 @@
 def jax_random_split(prng_key, num: int = 2):
     """Get a new key with ``jax.random.split``."""
     if prng_key is None:
-        return [None] * num
+        return (None,) * num
     # pylint: disable=import-outside-toplevel
     from jax.random import split
 
@@ -213,15 +213,11 @@ def measure_with_samples(
     """
     # last N measurements are sampling MCMs in ``dynamic_one_shot`` execution mode
     mps = measurements[0 : -len(mid_measurements)] if mid_measurements else measurements
-    skip_measure = any(v == -1 for v in mid_measurements.values()) if mid_measurements else False
 
     groups, indices = _group_measurements(mps)
 
     all_res = []
     for group in groups:
-        if skip_measure:
-            all_res.extend([None] * len(group))
-            continue
         if isinstance(group[0], ExpectationMP) and isinstance(
             group[0].obs, (Hamiltonian, LinearCombination)
         ):
@@ -477,11 +473,10 @@ def sample_state(
     # probabilities must be renormalized as they may not sum to one
     # see https://github.com/PennyLaneAI/pennylane/issues/5444
     norm = qml.math.sum(probs, axis=-1)
-    abs_diff = np.abs(norm - 1.0)
+    abs_diff = qml.math.abs(norm - 1.0)
     cutoff = 1e-07
 
     if is_state_batched:
-
         normalize_condition = False
 
         for s in abs_diff:
@@ -497,9 +492,9 @@ def sample_state(
         # rng.choice doesn't support broadcasting
         samples = np.stack([rng.choice(basis_states, shots, p=p) for p in probs])
     else:
-
-        if 0 < abs_diff < cutoff:
-            probs /= norm
+        if not 0 < abs_diff < cutoff:
+            norm = 1.0
+        probs = probs / norm
 
         samples = rng.choice(basis_states, shots, p=probs)
 

pennylane/devices/qubit/simulate.py

@@ -279,9 +279,33 @@ def simulate(
 
     has_mcm = any(isinstance(op, MidMeasureMP) for op in circuit.operations)
     if circuit.shots and has_mcm:
-        return simulate_one_shot_native_mcm(
-            circuit, debugger=debugger, rng=rng, prng_key=prng_key, interface=interface
+        results = []
+        aux_circ = qml.tape.QuantumScript(
+            circuit.operations,
+            circuit.measurements,
+            shots=[1],
+            trainable_params=circuit.trainable_params,
         )
+        keys = jax_random_split(prng_key, num=circuit.shots.total_shots)
+        if qml.math.get_deep_interface(circuit.data) == "jax":
+            # pylint: disable=import-outside-toplevel
+            import jax
+
+            def vv(k):
+                return simulate_one_shot_native_mcm(
+                    aux_circ, debugger=debugger, rng=rng, prng_key=k, interface=interface
+                )
+
+            results = jax.vmap(vv, in_axes=(0,))(keys)
+            results = list(zip(*results))
+        else:
+            for i, _ in enumerate(circuit.shots):
+                results.append(
+                    simulate_one_shot_native_mcm(
+                        aux_circ, debugger=debugger, rng=rng, prng_key=keys[i], interface=interface
+                    )
+                )
+        return tuple(results)
 
     ops_key, meas_key = jax_random_split(prng_key)
     state, is_state_batched = get_final_state(

pennylane/measurements/mid_measure.py

@@ -443,10 +443,10 @@ def __ge__(self, other):
         return self._transform_bin_op(lambda a, b: a >= b, other)
 
     def __and__(self, other):
-        return self._transform_bin_op(lambda a, b: a and b, other)
+        return self._transform_bin_op(qml.math.logical_and, other)
 
     def __or__(self, other):
-        return self._transform_bin_op(lambda a, b: a or b, other)
+        return self._transform_bin_op(qml.math.logical_or, other)
 
     def _apply(self, fn):
         """Apply a post computation to this measurement"""