[BUG] qml.measure can not be used within the body of qml.cond

[dwierichs]

Expected behavior

Any valid quantum subcircuit can be used within the body of qml.cond.
The resulting circuit is drawn correctly.

Actual behavior

Using qml.measure in the body of qml.cond raises an error.
Using qml.cond in the body of qml.cond works but drawing the circuit raises an error.

Additional information

No response

Source code

def sub_qfunc():
    mcm = qml.measure(0)

@qml.qnode(qml.device("default.qubit"))
def circuit():
    mcm = qml.measure(1)
    qml.cond(mcm, sub_qfunc)()
    return qml.expval(qml.Y(0) @ qml.Z(1))
circuit()

### Drawing example

def sub_qfunc(mcm):
    qml.cond(mcm, qml.X)(2)

@qml.qnode(qml.device("default.qubit"))
def circuit():
    mcm = qml.measure(0)
    mcm1 = qml.measure(1)
    qml.cond(mcm, sub_qfunc)(mcm1)
    return qml.expval(qml.Y(0) @ qml.Z(1))
qml.draw(circuit)()

Tracebacks

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[13], line 9
      7     qml.cond(mcm, sub_qfunc)()
      8     return qml.expval(qml.Y(0) @ qml.Z(1))
----> 9 circuit()

File ~/repos/pennylane/pennylane/workflow/qnode.py:1164, in QNode.__call__(self, *args, **kwargs)
   1162 if qml.capture.enabled():
   1163     return qml.capture.qnode_call(self, *args, **kwargs)
-> 1164 return self._impl_call(*args, **kwargs)

File ~/repos/pennylane/pennylane/workflow/qnode.py:1144, in QNode._impl_call(self, *args, **kwargs)
   1141     override_shots = kwargs["shots"]
   1143 # construct the tape
-> 1144 self.construct(args, kwargs)
   1146 original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device]
   1147 self._update_gradient_fn(shots=override_shots, tape=self._tape)

File ~/repos/pennylane/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File ~/repos/pennylane/pennylane/workflow/qnode.py:966, in QNode.construct(self, args, kwargs)
    964 with pldb_device_manager(self.device) as _:
    965     with qml.queuing.AnnotatedQueue() as q:
--> 966         self._qfunc_output = self.func(*args, **kwargs)
    968 self._tape = QuantumScript.from_queue(q, shots)
    970 params = self.tape.get_parameters(trainable_only=False)

Cell In[13], line 7, in circuit()
      4 @qml.qnode(qml.device("default.qubit"))
      5 def circuit():
      6     mcm = qml.measure(1)
----> 7     qml.cond(mcm, sub_qfunc)()
      8     return qml.expval(qml.Y(0) @ qml.Z(1))

File ~/repos/pennylane/pennylane/ops/op_math/condition.py:413, in cond.<locals>.wrapper(*args, **kwargs)
    410     raise ConditionalTransformError(with_meas_err)
    412 for op in qscript.operations:
--> 413     Conditional(condition, op)
    415 if false_fn is not None:
    416     # 2. Apply false_fn conditionally
    417     else_qscript = make_qscript(false_fn)(*args, **kwargs)

File ~/repos/pennylane/pennylane/capture/capture_meta.py:87, in CaptureMeta.__call__(cls, *args, **kwargs)
     83 if enabled():
     84     # when tracing is enabled, we want to
     85     # use bind to construct the class if we want class construction to add it to the jaxpr
     86     return cls._primitive_bind_call(*args, **kwargs)
---> 87 return type.__call__(cls, *args, **kwargs)

File ~/repos/pennylane/pennylane/ops/op_math/condition.py:58, in Conditional.__init__(self, expr, then_op, id)
     56 super().__init__(then_op, id=id)
     57 if self.grad_recipe is None:
---> 58     self.grad_recipe = [None] * self.num_params

File ~/repos/pennylane/pennylane/ops/op_math/condition.py:70, in Conditional.num_params(self)
     68 @property
     69 def num_params(self):
---> 70     return self.base.num_params

AttributeError: 'MidMeasureMP' object has no attribute 'num_params'


#######   Drawing output of second circuit:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[16], line 10
      8     qml.cond(mcm, sub_qfunc)(mcm1)
      9     return qml.expval(qml.Y(0) @ qml.Z(1))
---> 10 qml.draw(circuit)()

File ~/repos/pennylane/pennylane/drawer/draw.py:358, in _draw_qnode.<locals>.wrapper(*args, **kwargs)
    356 if tapes is not None:
    357     cache = {"tape_offset": 0, "matrices": []}
--> 358     res = [
    359         tape_text(
    360             t,
    361             wire_order=_wire_order,
    362             show_all_wires=show_all_wires,
    363             decimals=decimals,
    364             show_matrices=False,
    365             max_length=max_length,
    366             cache=cache,
    367         )
    368         for t in tapes
    369     ]
    370     if show_matrices and cache["matrices"]:
    371         mat_str = ""

File ~/repos/pennylane/pennylane/drawer/draw.py:359, in <listcomp>(.0)
    356 if tapes is not None:
    357     cache = {"tape_offset": 0, "matrices": []}
    358     res = [
--> 359         tape_text(
    360             t,
    361             wire_order=_wire_order,
    362             show_all_wires=show_all_wires,
    363             decimals=decimals,
    364             show_matrices=False,
    365             max_length=max_length,
    366             cache=cache,
    367         )
    368         for t in tapes
    369     ]
    370     if show_matrices and cache["matrices"]:
    371         mat_str = ""

File ~/repos/pennylane/pennylane/drawer/tape_text.py:501, in tape_text(tape, wire_order, show_all_wires, decimals, max_length, show_matrices, cache)
    499         tape_cache.append(op)
    500     else:
--> 501         layer_str = add_fn(op, layer_str, config)
    503 #################################################
    504 # Left justify layer strings and pad on the right
    505 #################################################
    506 # Adjust width for wire filler on unused wires
    507 max_label_len = max(len(s) for s in layer_str)

File ~/repos/pennylane/pennylane/drawer/tape_text.py:134, in _add_op(op, layer_str, config)
    132 if isinstance(op, qml.ops.Conditional):  # pylint: disable=no-member
    133     layer_str = _add_cond_grouping_symbols(op, layer_str, config)
--> 134     return _add_op(op.base, layer_str, config)
    136 if isinstance(op, MidMeasureMP):
    137     return _add_mid_measure_op(op, layer_str, config)

File ~/repos/pennylane/pennylane/drawer/tape_text.py:133, in _add_op(op, layer_str, config)
    131 """Updates ``layer_str`` with ``op`` operation."""
    132 if isinstance(op, qml.ops.Conditional):  # pylint: disable=no-member
--> 133     layer_str = _add_cond_grouping_symbols(op, layer_str, config)
    134     return _add_op(op.base, layer_str, config)
    136 if isinstance(op, MidMeasureMP):

File ~/repos/pennylane/pennylane/drawer/tape_text.py:88, in _add_cond_grouping_symbols(op, layer_str, config)
     85 n_wires = len(config.wire_map)
     87 mapped_wires = [config.wire_map[w] for w in op.wires]
---> 88 mapped_bits = [config.bit_map[m] for m in op.meas_val.measurements]
     89 max_w = max(mapped_wires)
     90 max_b = max(mapped_bits) + n_wires

File ~/repos/pennylane/pennylane/drawer/tape_text.py:88, in <listcomp>(.0)
     85 n_wires = len(config.wire_map)
     87 mapped_wires = [config.wire_map[w] for w in op.wires]
---> 88 mapped_bits = [config.bit_map[m] for m in op.meas_val.measurements]
     89 max_w = max(mapped_wires)
     90 max_b = max(mapped_bits) + n_wires

KeyError: measure(wires=[1])

System information

pl dev

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.