[BUG] qml.draw, qml.draw_mpl do no work with qml.defer_measurements when collecting MCM statistics

[mudit2812]

Expected behavior

The drawer correctly draws the circuit with deferred measurements

Actual behavior

An error is raised

Additional information

No response

Source code

import pennylane as qml
import numpy as np

@qml.qnode(qml.device("default.qubit"))
def node(x):
    qml.RX(x, 0)
    mcm = qml.measure(0, postselect=1)
    qml.RZ(0.5, 0)
    return qml.expval(qml.Z(0)), qml.expval(op=mcm)

print(node(np.pi/4)) # Works
print(qml.defer_measurements(node)(np.pi/4)) # Works
print(qml.draw(node)(np.pi/4)) # Works
print(qml.draw_mpl(qml.defer_measurements(node))(np.pi/4)) # does not work

Tracebacks

(tensor(-1., requires_grad=True), tensor(1., requires_grad=True))
(tensor(-1., requires_grad=True), tensor(1., requires_grad=True))
0: ──RX(0.79)──┤↗₁├──RZ(0.50)─┤  <Z>  
                ╚═════════════╡  <MCM>
{
	"name": "TypeError",
	"message": "list indices must be integers or slices, not NoneType",
	"stack": "---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[4], line 14
     12 print(qml.defer_measurements(node)(np.pi/4)) # Works
     13 print(qml.draw(node)(np.pi/4)) # Works
---> 14 print(qml.draw_mpl(qml.defer_measurements(node))(np.pi/4)) # does not work

File ~/repos/pennylane/pennylane/drawer/draw.py:630, in _draw_mpl_qnode.<locals>.wrapper(*args, **kwargs_qnode)
    627     except TypeError:
    628         _wire_order = tape.wires
--> 630 return tape_mpl(
    631     tape,
    632     wire_order=_wire_order,
    633     show_all_wires=show_all_wires,
    634     decimals=decimals,
    635     style=style,
    636     fig=fig,
    637     **kwargs,
    638 )

File ~/repos/pennylane/pennylane/drawer/tape_mpl.py:455, in tape_mpl(tape, wire_order, show_all_wires, decimals, style, fig, **kwargs)
    453     _set_style(style)
    454 try:
--> 455     return _tape_mpl(
    456         tape,
    457         wire_order=wire_order,
    458         show_all_wires=show_all_wires,
    459         decimals=decimals,
    460         fig=fig,
    461         **kwargs,
    462     )
    463 finally:
    464     if update_style:
    465         # we don't want to mess with how it modifies whether the interface is interactive
    466         # but we want to restore everything else

File ~/repos/pennylane/pennylane/drawer/tape_mpl.py:231, in _tape_mpl(tape, wire_order, show_all_wires, decimals, fig, **kwargs)
    228 n_layers = len(layers)
    229 n_wires = len(wire_map)
--> 231 cwire_layers, cwire_wires = cwire_connections(layers + [tape.measurements], bit_map)
    233 drawer = MPLDrawer(
    234     n_layers=n_layers,
    235     n_wires=n_wires,
   (...)
    238     fig=fig,
    239 )
    241 config = _Config(
    242     decimals=decimals,
    243     active_wire_notches=active_wire_notches,
    244     bit_map=bit_map,
    245     terminal_layers=[cl[-1] for cl in cwire_layers],
    246 )

File ~/repos/pennylane/pennylane/drawer/utils.py:197, in cwire_connections(layers, bit_map)
    195     for m in op.mv.measurements:
    196         cwire = bit_map[m]
--> 197         connected_layers[cwire].append(layer_idx)
    198 else:
    199     for m in op.mv:

TypeError: list indices must be integers or slices, not NoneType"
}

System information

Core dev

Existing GitHub issues

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