[Capture] Improve cond by storing missing false branches as empty jaxpr (#8080)

**Context:**

Right now we store missing false branches in a cond equation with
`None`. This PR simply stores them as jaxpr with no output instead. This
just means we have fewer logical branches we need to consider when
handling `cond` equations, which will make maintainence easier.

**Description of the Change:**

Store missing false branches as empty jaxpr instead of `None`.

**Benefits:**

A reduction in logical complexity.

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-97723]

Author: albi3ro

doc/releases/changelog-dev.md

@@ -453,6 +453,10 @@
 
 <h3>Internal changes ⚙️</h3>
 
+* The `cond` primitive with program capture no longer stores missing false branches as `None`, instead storing them
+  as jaxprs with no output.
+  [(#8080)](https://github.com/PennyLaneAI/pennylane/pull/8080)
+
 * Removed unnecessary execution tests along with accuracy validation in `tests/ops/functions/test_map_wires.py`.
   [(#8032)](https://github.com/PennyLaneAI/pennylane/pull/8032)
 

pennylane/capture/base_interpreter.py

@@ -533,15 +533,11 @@ def handle_cond(self, *invals, jaxpr_branches, consts_slices, args_slice):
 
     for const_slice, jaxpr in zip(consts_slices, jaxpr_branches):
         consts = invals[const_slice]
-        if jaxpr is None:
-            new_jaxprs.append(None)
-            new_consts_slices.append(slice(0, 0))
-        else:
-            new_jaxpr = jaxpr_to_jaxpr(copy(self), jaxpr, consts, *args)
-            new_jaxprs.append(new_jaxpr.jaxpr)
-            new_consts.extend(new_jaxpr.consts)
-            new_consts_slices.append(slice(end_const_ind, end_const_ind + len(new_jaxpr.consts)))
-            end_const_ind += len(new_jaxpr.consts)
+        new_jaxpr = jaxpr_to_jaxpr(copy(self), jaxpr, consts, *args)
+        new_jaxprs.append(new_jaxpr.jaxpr)
+        new_consts.extend(new_jaxpr.consts)
+        new_consts_slices.append(slice(end_const_ind, end_const_ind + len(new_jaxpr.consts)))
+        end_const_ind += len(new_jaxpr.consts)
 
     new_args_slice = slice(end_const_ind, None)
     return cond_prim.bind(
@@ -694,7 +690,7 @@ def flattened_cond(self, *invals, jaxpr_branches, consts_slices, args_slice):
             raise NotImplementedError(
                 f"{self} does not yet support jitting cond with abstract conditions."
             )
-        if pred and jaxpr is not None:
+        if pred:
             return copy(self).eval(jaxpr, consts, *args)
     return ()
 

pennylane/decomposition/collect_resource_ops.py

@@ -83,11 +83,10 @@ def explore_all_branches(self, *invals, jaxpr_branches, consts_slices, args_slic
     outvals = ()
     for _, jaxpr, consts_slice in zip(conditions, jaxpr_branches, consts_slices):
         consts = invals[consts_slice]
-        if jaxpr is not None:
-            dummy = copy(self).eval(jaxpr, consts, *args)
-            # The cond_prim may or may not expect outvals, so we need to check whether
-            # the first branch returns something significant. If so, we use the return
-            # value of the first branch as the outvals of this cond_prim.
-            if dummy and not outvals:
-                outvals = dummy
+        dummy = copy(self).eval(jaxpr, consts, *args)
+        # The cond_prim may or may not expect outvals, so we need to check whether
+        # the first branch returns something significant. If so, we use the return
+        # value of the first branch as the outvals of this cond_prim.
+        if dummy and not outvals:
+            outvals = dummy
     return outvals

pennylane/ops/op_math/condition.py

@@ -75,6 +75,10 @@ def new_fn(*args, **kwargs):
     return fn
 
 
+def _empty_return_fn(*_, **__):
+    return None
+
+
 class Conditional(SymbolicOp, Operation):
     """A Conditional Operation.
 
@@ -286,19 +290,17 @@ def __call_capture_enabled(self, *args, **kwargs):
                 raise ValueError(f"Condition predicate must be a scalar. Got {pred_shape}.")
             conditions.append(pred)
             if fn is None:
-                jaxpr_branches.append(None)
-                consts_slices.append(slice(0, 0))
-            else:
-                f = FlatFn(functools.partial(fn, **kwargs))
-                if jax.config.jax_dynamic_shapes:
-                    f = _add_abstract_shapes(f)
-                jaxpr = jax.make_jaxpr(f, abstracted_axes=abstracted_axes)(*args)
-                jaxpr_branches.append(jaxpr.jaxpr)
-                consts_slices.append(slice(end_const_ind, end_const_ind + len(jaxpr.consts)))
-                consts += jaxpr.consts
-                end_const_ind += len(jaxpr.consts)
-
-        _validate_jaxpr_returns(jaxpr_branches)
+                fn = _empty_return_fn
+            f = FlatFn(functools.partial(fn, **kwargs))
+            if jax.config.jax_dynamic_shapes:
+                f = _add_abstract_shapes(f)
+            jaxpr = jax.make_jaxpr(f, abstracted_axes=abstracted_axes)(*args)
+            jaxpr_branches.append(jaxpr.jaxpr)
+            consts_slices.append(slice(end_const_ind, end_const_ind + len(jaxpr.consts)))
+            consts += jaxpr.consts
+            end_const_ind += len(jaxpr.consts)
+
+        _validate_jaxpr_returns(jaxpr_branches, self.otherwise_fn)
         flat_args, _ = jax.tree_util.tree_flatten(args)
         results = cond_prim.bind(
             *conditions,
@@ -746,21 +748,15 @@ def _validate_abstract_values(
             _aval_mismatch_error(branch_type, branch_index, i, outval, expected_outval)
 
 
-def _validate_jaxpr_returns(jaxpr_branches):
+def _validate_jaxpr_returns(jaxpr_branches, false_fn):
     out_avals_true = [out.aval for out in jaxpr_branches[0].outvars]
-    for idx, jaxpr_branch in enumerate(jaxpr_branches):
-
-        if idx == 0:
-            continue
 
-        if jaxpr_branch is None:
-            if out_avals_true:
-                raise ValueError(
-                    "The false branch must be provided if the true branch returns any variables"
-                )
-            # this is tested, but coverage does not pick it up
-            continue  # pragma: no cover
+    if false_fn is None and out_avals_true:
+        raise ValueError(
+            "The false branch must be provided if the true branch returns any variables"
+        )
 
+    for idx, jaxpr_branch in enumerate(jaxpr_branches[1:], start=1):
         out_avals_branch = [out.aval for out in jaxpr_branch.outvars]
         branch_type = "elif" if idx < len(jaxpr_branches) - 1 else "false"
         _validate_abstract_values(out_avals_branch, out_avals_true, branch_type, idx - 1)
@@ -805,8 +801,6 @@ def _(*all_args, jaxpr_branches, consts_slices, args_slice):
 
         for pred, jaxpr, const_slice in zip(conditions, jaxpr_branches, consts_slices):
             consts = all_args[const_slice]
-            if jaxpr is None:
-                continue
             if isinstance(pred, qml.measurements.MeasurementValue):
 
                 with qml.queuing.AnnotatedQueue() as q:

pennylane/tape/plxpr_conversion.py

@@ -150,8 +150,6 @@ def _(self, *all_args, jaxpr_branches, consts_slices, args_slice):
 
     for pred, jaxpr, const_slice in zip(conditions, jaxpr_branches, consts_slices):
         consts = all_args[const_slice]
-        if jaxpr is None:
-            continue
         if isinstance(pred, qml.measurements.MeasurementValue):
             if jaxpr.outvars:
                 outvals = [v.aval for v in jaxpr.outvars]

pennylane/transforms/defer_measurements.py

@@ -408,11 +408,6 @@ def _(self, *invals, jaxpr_branches, consts_slices, args_slice):
         args = invals[args_slice]
 
         for i, (condition, jaxpr) in enumerate(zip(conditions, jaxpr_branches, strict=True)):
-            if jaxpr is None:
-                # If a false branch isn't provided, the jaxpr corresponding to the condition
-                # for the false branch will be None. That is the only scenario where we would
-                # reach here.
-                continue
 
             if isinstance(condition, MeasurementValue):
                 control_wires = Wires([m.wires[0] for m in condition.measurements])

pennylane/transforms/optimization/merge_amplitude_embedding.py

@@ -251,33 +251,27 @@ def _(self, *invals, jaxpr_branches, consts_slices, args_slice):
 
         for const_slice, jaxpr in zip(consts_slices, jaxpr_branches):
             consts = invals[const_slice]
-            if jaxpr is None:
-                new_jaxprs.append(None)
-                new_consts_slices.append(slice(0, 0))
-            else:
-                new_jaxpr = jaxpr_to_jaxpr(copy(self), jaxpr, consts, *args)
-
-                # Update state so far so collisions with
-                # newly seen states from the branches continue to be
-                # detected after the cond
-                curr_wires |= self.state["visited_wires"]
-                curr_dynamic_wires_found = self.state["dynamic_wires_found"]
-                curr_ops_found = self.state["ops_found"]
-
-                # Reset state for the next branch so we don't get false positive collisions
-                # (copy so if state mutates we preserved true initial state)
-                self.state = {
-                    "visited_wires": copy(initial_wires),
-                    "dynamic_wires_found": initial_dynamic_wires_found,
-                    "ops_found": initial_ops_found,
-                }
-
-                new_jaxprs.append(new_jaxpr.jaxpr)
-                new_consts.extend(new_jaxpr.consts)
-                new_consts_slices.append(
-                    slice(end_const_ind, end_const_ind + len(new_jaxpr.consts))
-                )
-                end_const_ind += len(new_jaxpr.consts)
+            new_jaxpr = jaxpr_to_jaxpr(copy(self), jaxpr, consts, *args)
+
+            # Update state so far so collisions with
+            # newly seen states from the branches continue to be
+            # detected after the cond
+            curr_wires |= self.state["visited_wires"]
+            curr_dynamic_wires_found = curr_dynamic_wires_found or self.state["dynamic_wires_found"]
+            curr_ops_found = curr_ops_found or self.state["ops_found"]
+
+            # Reset state for the next branch so we don't get false positive collisions
+            # (copy so if state mutates we preserved true initial state)
+            self.state = {
+                "visited_wires": copy(initial_wires),
+                "dynamic_wires_found": initial_dynamic_wires_found,
+                "ops_found": initial_ops_found,
+            }
+
+            new_jaxprs.append(new_jaxpr.jaxpr)
+            new_consts.extend(new_jaxpr.consts)
+            new_consts_slices.append(slice(end_const_ind, end_const_ind + len(new_jaxpr.consts)))
+            end_const_ind += len(new_jaxpr.consts)
 
         # Reset state to all updates from all branches in the cond
         self.state = {

tests/capture/test_base_interpreter.py

@@ -501,7 +501,8 @@ def f():
 
         jaxpr = jax.make_jaxpr(f)(True)
 
-        assert jaxpr.eqns[0].params["jaxpr_branches"][-1] is None  # no false branch
+        false_branch = jaxpr.eqns[0].params["jaxpr_branches"][-1]
+        assert len(false_branch.eqns) == 0
 
         with qml.queuing.AnnotatedQueue() as q_true:
             jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, True)

tests/capture/test_capture_cond.py

@@ -409,7 +409,9 @@ def f():
         assert len(true_fn.outvars) == 0
         assert true_fn.eqns[0].primitive == qml.X._primitive  # pylint: disable=protected-access
 
-        assert jaxpr.eqns[0].params["jaxpr_branches"][-1] is None
+        false_fn = jaxpr.eqns[0].params["jaxpr_branches"][-1]
+        assert len(false_fn.eqns) == 0
+        assert len(false_fn.outvars) == 0
 
 
 dev = qml.device("default.qubit", wires=3)