Repository: PennyLaneAI/pennylane
Pull request #4930 (2023-12-08 15:24:58)

Support drawing mid-circuit measurement statistics with `qml.draw`
**Context:**
This PR adds support for drawing the collection of mid-circuit measurement statistics with `qml.draw`. An example is shown below:
```python
def circ():
    qml.Hadamard(0)
    m0 = qml.measure(0)
    qml.Hadamard(1)
    m1 = qml.measure(1)
    qml.Hadamard(2)
    m2 = qml.measure(2)
    qml.Hadamard(3)
    qml.measure(3)
    qml.cond(m0, qml.PauliX)(0)
    qml.cond(m0 & m1, qml.PauliY)(1)
    return qml.expval(m2), qml.sample([m1, m2])

print(qml.draw(circ)())
```
```
0: ──H──┤↗├──────────────────────────X────┤                    
1: ──────║───H──┤↗├──────────────────║──Y─┤                    
2: ──────║───────║───H──┤↗├──────────║──║─┤                    
3: ──────║───────║───────║───H──┤↗├──║──║─┤                    
         ╚═══════║═══════║═══════════╩══╣                      
                 ╚═══════║══════════════╩═╡        ╭Sample[MCM]
                         ╚════════════════╡  <MCM> ╰Sample[MCM]
```

**Description of the Change:**

* Refactor `tape_text` to only run through one loop containing all layers, not separate loops for op layers and measurement layers. The logic inside the loop has been adjusted accordingly. I also added big block comments to make the separation between different steps more obvious.
* Add helpers to `tape_text.py` to add grouping symbols and labels for mid-circuit measurement statistics correctly.
* Update `drawable_layers` to find layers for mid-circuit measurement statistics correctly. This included adding a new helper to recursively find layers for terminal measurements that use measurement values.
* The bit map is now initialized with `None` values before everything else, and this is used to keep track of which mid-circuit measurements are used. Mostly a refactor to tidy up the code.
* Remove code from `draw_mpl` that deferred measurements unconditionally for qnodes with mid-circuit measurements. This code was removed from `qml.draw` in #4901 and `qml.draw_mpl` is now updated as well.
* The example in the `iterative_qpe` docstring was written before we had advanced mid-circuit measurement support and drawing capabilities, so I updated the example as well.

**Benefits:**
* We can draw mid-circuit measurement statistics!
* Code related to the text drawer is slightly more readable.

**Possible Drawbacks:**

**Related GitHub Issues:**


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4930?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`dcd1201`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/dcd120100872f594817268cb62e4af224008bde1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66% compared to head [(`ca35577`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4930?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66%.


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff            @@
##           master    #4930    +/-   ##
========================================
  Coverage   99.66%   99.66%            
========================================
  Files         388      388            
  Lines       35285    35019   -266     
========================================
- Hits        35166    34901   -265     
+ Misses        119      118     -1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4930?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

mudit2812: [sc-49830]

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 401db9c52..f8c1fabd4 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -14,9 +14,10 @@
   [(#4901)](https://github.com/PennyLaneAI/pennylane/pull/4901)
   [(#4850)](https://github.com/PennyLaneAI/pennylane/pull/4850)
   [(#4917)](https://github.com/PennyLaneAI/pennylane/pull/4917)
+  [(#4930)](https://github.com/PennyLaneAI/pennylane/pull/4930)
 
   Drawing of mid-circuit measurement capabilities including qubit reuse and reset,
-  postselection, and conditioning are supported.
+  postselection, conditioning, and collecting statistics is supported.
 
   ```python
   import pennylane as qml
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index 10337de0d..66f0a7cf3 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -573,16 +573,7 @@ def _draw_mpl_qnode(
                 qnode.construct(args, kwargs_qnode)
                 if isinstance(qnode.device, qml.devices.Device):
                     program = qnode.transform_program
-                    if any(
-                        isinstance(op, qml.measurements.MidMeasureMP)
-                        for op in qnode.tape.operations
-                    ):
-                        tapes, _ = qml.defer_measurements(qnode.tape, device=qnode.device)
-                    else:
-                        tapes = [qnode.tape]
-
-                    tapes, _ = program(tapes)
-                    tape = tapes[0]
+                    [tape], _ = program([qnode.tape])
                 else:
                     tape = qnode.tape
             finally:
diff --git a/pennylane/drawer/drawable_layers.py b/pennylane/drawer/drawable_layers.py
index cc8047e81..8ffc0e69b 100644
--- a/pennylane/drawer/drawable_layers.py
+++ b/pennylane/drawer/drawable_layers.py
@@ -16,7 +16,7 @@ This module contains a helper function to sort operations into layers.
 """
 
 from pennylane.ops import Conditional
-from pennylane.measurements import MidMeasureMP
+from pennylane.measurements import MidMeasureMP, MeasurementProcess
 from .utils import default_wire_map
 
 
@@ -50,12 +50,43 @@ def _recursive_find_layer(layer_to_check, op_occupied_wires, occupied_wires_per_
     return _recursive_find_layer(layer_to_check - 1, op_occupied_wires, occupied_wires_per_layer)
 
 
-def _get_op_occupied_wires(op, wire_map, cond_measurements):
+def _recursive_find_mcm_stats_layer(
+    layer_to_check, op_occupied_cwires, used_cwires_per_layer, bit_map
+):
+    """Determine correct layer for a terminal measurement that is collectings statistics
+    for mid-circuit measurement values.
+
+    Args:
+        layer_to_check (int): the function determines if the operation fits on this layer
+        op_occupied_cwires (set(int)): classical wires occupied by measurement
+        occupied_cwires_per_layer (list[set[int]]): which classical wires are already
+            in use for collecting statistics. Each set is a different layer.
+        bit_map (dict): A map containing mid-circuit measurements used for classical conditions
+            or collecting statistics as keys
+
+    Returns:
+        int: layer to place measurement process in
+    """
+
+    if op_occupied_cwires & used_cwires_per_layer[layer_to_check]:
+        # this layer is occupied, use higher one
+        return layer_to_check + 1
+
+    if layer_to_check == 0:
+        # reached first layer, so stop
+        return 0
+    # keep pushing the operation to lower layers
+    return _recursive_find_mcm_stats_layer(
+        layer_to_check - 1, op_occupied_cwires, used_cwires_per_layer, bit_map
+    )
+
+
+def _get_op_occupied_wires(op, wire_map, bit_map):
     """Helper function to find wires that would be used by an operator in a drawable layer."""
     if isinstance(op, MidMeasureMP):
         mapped_wire = wire_map[op.wires[0]]
 
-        if op in cond_measurements:
+        if op in bit_map:
             min_wire = mapped_wire
             max_wire = max(wire_map.values())
             return set(range(min_wire, max_wire + 1))
@@ -82,18 +113,21 @@ def _get_op_occupied_wires(op, wire_map, cond_measurements):
     return set(range(min_wire, max_wire + 1))
 
 
-def drawable_layers(ops, wire_map=None):
+def drawable_layers(operations, wire_map=None, bit_map=None):
     """Determine non-overlapping yet dense placement of operations into layers for drawing.
 
     Args:
-        ops Iterable[~.Operator]: a list of operations
+        operations (Iterable[~.Operator]): A list of operations.
 
     Keyword Args:
-        wire_map=None (dict): a map from wire label to non-negative integers
+        wire_map (dict): A map from wire label to non-negative integers. Defaults to None.
+        bit_map (dict): A map containing mid-circuit measurements used for classical conditions
+            or collecting statistics as keys. Defaults to None.
 
     Returns:
-        list[set[~.Operator]] : Each index is a set of operations
-        for the corresponding layer
+        (list[set[~.Operator]], list[set[~.MeasurementProcess]]) : Each index is a set of operations
+        for the corresponding layer in both lists. The first list corresponds to the operation layers,
+        and the second corresponds to the measurement layers.
 
     **Details**
 
@@ -115,41 +149,55 @@ def drawable_layers(ops, wire_map=None):
     transformation.
 
     """
-    # pylint:disable=too-many-nested-blocks,too-many-branches,too-many-statements
 
-    if not wire_map:
-        wire_map = default_wire_map(ops)
-    # initialize
+    wire_map = wire_map or default_wire_map(operations)
+    bit_map = bit_map or {}
+
+    # initialize for operation layers
     max_layer = 0
     occupied_wires_per_layer = [set()]
     ops_in_layer = [[]]
-    measured_wires = {}
-
-    # Collect all mid-circuit measurements used for classical conditioning
-    cond_measurements = set()
-    for op in ops:
-        if isinstance(op, Conditional):
-            cond_measurements.update(op.meas_val.measurements)
+    used_cwires_per_layer = [set()]
 
     # loop over operations
-    for op in ops:
+    for op in operations:
         if isinstance(op, MidMeasureMP):
             if len(op.wires) > 1:
                 raise ValueError("Cannot draw mid-circuit measurements with more than one wire.")
 
-            measured_wires[op.id] = wire_map[op.wires[0]]
-
-        op_occupied_wires = _get_op_occupied_wires(op, wire_map, cond_measurements)
-        op_layer = _recursive_find_layer(max_layer, op_occupied_wires, occupied_wires_per_layer)
+        if isinstance(op, MeasurementProcess) and op.mv is not None:
+            # Only terminal measurements that collect mid-circuit measurement statistics have
+            # op.mv != None.
+            # Get the occupied classical wires of the measurement process and find which layer
+            # to put it in.
+            op_occupied_wires = set()
+            mapped_cwires = (
+                [bit_map[m.measurements[0]] for m in op.mv]
+                if isinstance(op.mv, list)
+                else [bit_map[m] for m in op.mv.measurements]
+            )
+            op_occupied_cwires = set(range(min(mapped_cwires), max(mapped_cwires) + 1))
+            op_layer = _recursive_find_mcm_stats_layer(
+                max_layer, op_occupied_cwires, used_cwires_per_layer, bit_map
+            )
+
+        else:
+            # Find occupied wires of the operator/measurement process and find which layer to
+            # put it in.
+            op_occupied_wires = _get_op_occupied_wires(op, wire_map, bit_map)
+            op_layer = _recursive_find_layer(max_layer, op_occupied_wires, occupied_wires_per_layer)
+            op_occupied_cwires = set()
 
         # see if need to add new layer
         if op_layer > max_layer:
             max_layer += 1
-            occupied_wires_per_layer.append(set(measured_wires.values()))
+            occupied_wires_per_layer.append(set())
             ops_in_layer.append([])
+            used_cwires_per_layer.append(set())
 
         # add to op_layer
         ops_in_layer[op_layer].append(op)
         occupied_wires_per_layer[op_layer].update(op_occupied_wires)
+        used_cwires_per_layer[op_layer].update(op_occupied_cwires)
 
     return list(filter(None, ops_in_layer[:-1])) + ops_in_layer[-1:]
diff --git a/pennylane/drawer/tape_mpl.py b/pennylane/drawer/tape_mpl.py
index 58782d987..463e8241d 100644
--- a/pennylane/drawer/tape_mpl.py
+++ b/pennylane/drawer/tape_mpl.py
@@ -27,7 +27,7 @@ from pennylane import ops
 from pennylane.measurements import MidMeasureMP
 from .mpldrawer import MPLDrawer
 from .drawable_layers import drawable_layers
-from .utils import convert_wire_order, unwrap_controls, cwire_connections
+from .utils import convert_wire_order, unwrap_controls, cwire_connections, default_bit_map
 from .style import _set_style
 
 has_mpl = True
@@ -200,8 +200,9 @@ def _tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, **kwar
 
     wire_map = convert_wire_order(tape, wire_order=wire_order, show_all_wires=show_all_wires)
     tape = qml.map_wires(tape, wire_map=wire_map)[0][0]
+    bit_map = default_bit_map(tape)
 
-    layers = drawable_layers(tape.operations, {i: i for i in tape.wires})
+    layers = drawable_layers(tape.operations, wire_map={i: i for i in tape.wires}, bit_map=bit_map)
 
     for i, layer in enumerate(layers):
         if any(isinstance(o, qml.measurements.MidMeasureMP) and o.reset for o in layer):
@@ -210,7 +211,7 @@ def _tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, **kwar
     n_layers = len(layers)
     n_wires = len(wire_map)
 
-    bit_map, cwire_layers, cwire_wires = cwire_connections(layers)
+    cwire_layers, cwire_wires = cwire_connections(layers, bit_map)
 
     drawer = MPLDrawer(
         n_layers=n_layers, n_wires=n_wires, c_wires=len(bit_map), wire_options=wire_options
diff --git a/pennylane/drawer/tape_text.py b/pennylane/drawer/tape_text.py
index 76afa629f..890a5bb5f 100644
--- a/pennylane/drawer/tape_text.py
+++ b/pennylane/drawer/tape_text.py
@@ -20,10 +20,18 @@ This module contains logic for the text based circuit drawer through the ``tape_
 from dataclasses import dataclass
 from typing import Optional
 import pennylane as qml
-from pennylane.measurements import Expectation, Probability, Sample, Variance, State, MidMeasureMP
+from pennylane.measurements import (
+    Expectation,
+    Probability,
+    Sample,
+    Variance,
+    State,
+    Counts,
+    MidMeasureMP,
+)
 
 from .drawable_layers import drawable_layers
-from .utils import convert_wire_order, unwrap_controls, cwire_connections
+from .utils import convert_wire_order, unwrap_controls, cwire_connections, default_bit_map
 
 
 @dataclass
@@ -160,13 +168,51 @@ measurement_label_map = {
     Expectation: lambda label: f"<{label}>",
     Probability: lambda label: f"Probs[{label}]" if label else "Probs",
     Sample: lambda label: f"Sample[{label}]" if label else "Sample",
+    Counts: lambda label: f"Counts[{label}]" if label else "Counts",
     Variance: lambda label: f"Var[{label}]",
     State: lambda label: "State",
 }
 
 
+def _add_cwire_measurement_grouping_symbols(mcms, layer_str, config):
+    """Adds symbols indicating the extent of a given object for mid-circuit measurement
+    statistics."""
+    if len(mcms) > 1:
+        n_wires = len(config.wire_map)
+        mapped_bits = [config.bit_map[m] for m in mcms]
+        min_b, max_b = min(mapped_bits) + n_wires, max(mapped_bits) + n_wires
+
+        layer_str[min_b] = "╭"
+        layer_str[max_b] = "╰"
+
+        for b in range(min_b + 1, max_b):
+            layer_str[b] = "├" if b - n_wires in mapped_bits else "│"
+
+    return layer_str
+
+
+def _add_cwire_measurement(m, layer_str, config):
+    """Updates ``layer_str`` with the ``m`` measurement when it is used
+    for collecting mid-circuit measurement statistics."""
+    mcms = [v.measurements[0] for v in m.mv] if isinstance(m.mv, list) else m.mv.measurements
+    layer_str = _add_cwire_measurement_grouping_symbols(mcms, layer_str, config)
+
+    mv_label = "MCM"
+    meas_label = measurement_label_map[m.return_type](mv_label)
+
+    n_wires = len(config.wire_map)
+    for mcm in mcms:
+        ind = config.bit_map[mcm] + n_wires
+        layer_str[ind] += meas_label
+
+    return layer_str
+
+
 def _add_measurement(m, layer_str, config):
     """Updates ``layer_str`` with the ``m`` measurement."""
+    if m.mv is not None:
+        return _add_cwire_measurement(m, layer_str, config)
+
     layer_str = _add_grouping_symbols(m, layer_str, config)
 
     if m.obs is None:
@@ -383,24 +429,22 @@ def tape_text(
     tape_cache = []
 
     wire_map = convert_wire_order(tape, wire_order=wire_order, show_all_wires=show_all_wires)
+    bit_map = default_bit_map(tape)
     n_wires = len(wire_map)
+    n_bits = len(bit_map)
     if n_wires == 0:
         return ""
 
     # Used to store lines that are hitting the maximum length
     finished_lines = []
 
-    layers_list = [
-        drawable_layers(tape.operations, wire_map=wire_map),
-        drawable_layers(tape.measurements, wire_map=wire_map),
-    ]
-    add_list = [_add_op, _add_measurement]
-    wire_fillers = ["─", " "]
-    bit_fillers = ["═", " "]
-    enders = [True, False]  # add "─┤" after all operations
+    layers = drawable_layers(tape.operations, wire_map=wire_map, bit_map=bit_map)
+    final_operations_layer = len(layers) - 1
+    layers += drawable_layers(tape.measurements, wire_map=wire_map, bit_map=bit_map)
 
-    bit_map, cwire_layers, _ = cwire_connections(layers_list[0] + layers_list[1])
-    n_bits = len(bit_map)
+    # Update bit map and collect information about connections between mid-circuit measurements,
+    # classical conditions, and terminal measurements for processing mid-circuit measurements.
+    cwire_layers, _ = cwire_connections(layers, bit_map)
 
     wire_totals = [f"{wire}: " for wire in wire_map]
     bit_totals = ["" for _ in range(n_bits)]
@@ -409,93 +453,96 @@ def tape_text(
     wire_totals = [s.rjust(line_length, " ") for s in wire_totals]
     bit_totals = [s.rjust(line_length, " ") for s in bit_totals]
 
-    for layers, add, w_filler, b_filler, ender in zip(
-        layers_list, add_list, wire_fillers, bit_fillers, enders
-    ):
-        # Collect information needed for drawing layers
-        config = _Config(
-            wire_map=wire_map,
-            bit_map=bit_map,
-            cur_layer=-1,
-            cwire_layers=cwire_layers,
-            decimals=decimals,
-            cache=cache,
-        )
-
-        for i, layer in enumerate(layers):
-            # Add filler before current layer
-            layer_str = [w_filler] * n_wires + [" "] * n_bits
-            for b in bit_map.values():
-                cur_b_filler = b_filler if min(cwire_layers[b]) < i < max(cwire_layers[b]) else " "
-                layer_str[b + n_wires] = cur_b_filler
-
-            config.cur_layer = i
-            # Keep track of mid-circuit measurements in each layer that are used
-            # for conditions, if any
-            cur_layer_mid_measure = None
-
-            for op in layer:
-                if isinstance(op, qml.tape.QuantumScript):
-                    layer_str = _add_grouping_symbols(op, layer_str, config)
-                    label = f"Tape:{cache['tape_offset']+len(tape_cache)}"
-                    for w in op.wires:
-                        layer_str[wire_map[w]] += label
-                    tape_cache.append(op)
-                else:
-                    layer_str = add(op, layer_str, config)
-
-                    if isinstance(op, MidMeasureMP) and op in bit_map:
-                        cur_layer_mid_measure = op
-
-            # Adjust width for wire filler on unused wires
-            max_label_len = max(len(s) for s in layer_str)
-            for w in range(n_wires):
-                layer_str[w] = layer_str[w].ljust(max_label_len, w_filler)
-
-            # Adjust width for bit filler on unused bits
+    # Collect information needed for drawing layers
+    config = _Config(
+        wire_map=wire_map,
+        bit_map=bit_map,
+        cur_layer=-1,
+        cwire_layers=cwire_layers,
+        decimals=decimals,
+        cache=cache,
+    )
+
+    for i, layer in enumerate(layers):
+        # Update fillers and helper function
+        w_filler = "─" if i <= final_operations_layer else " "
+        b_filler = "═" if i <= final_operations_layer else " "
+        add_fn = _add_op if i <= final_operations_layer else _add_measurement
+
+        # Create initial strings for the current layer using wire and cwire fillers
+        layer_str = [w_filler] * n_wires + [" "] * n_bits
+        for b in bit_map.values():
+            cur_b_filler = b_filler if min(cwire_layers[b]) < i < max(cwire_layers[b]) else " "
+            layer_str[b + n_wires] = cur_b_filler
+
+        config.cur_layer = i
+
+        ##########################################
+        # Update current layer strings with labels
+        ##########################################
+        for op in layer:
+            if isinstance(op, qml.tape.QuantumScript):
+                layer_str = _add_grouping_symbols(op, layer_str, config)
+                label = f"Tape:{cache['tape_offset']+len(tape_cache)}"
+                for w in op.wires:
+                    layer_str[wire_map[w]] += label
+                tape_cache.append(op)
+            else:
+                layer_str = add_fn(op, layer_str, config)
+
+        #################################################
+        # Left justify layer strings and pad on the right
+        #################################################
+        # Adjust width for wire filler on unused wires
+        max_label_len = max(len(s) for s in layer_str)
+        for w in range(n_wires):
+            layer_str[w] = layer_str[w].ljust(max_label_len, w_filler)
+
+        # Adjust width for bit filler on unused bits
+        for b in range(n_bits):
+            cur_b_filler = b_filler if cwire_layers[b][0] <= i < cwire_layers[b][-1] else " "
+            layer_str[b + n_wires] = layer_str[b + n_wires].ljust(max_label_len, cur_b_filler)
+
+        line_length += max_label_len + 1  # one for the filler character
+
+        ##################
+        # Create new lines
+        ##################
+        if line_length > max_length:
+            # move totals into finished_lines and reset totals
+            finished_lines += wire_totals + bit_totals
+            finished_lines[-1] += "\n"
+            wire_totals = [w_filler] * n_wires
+
+            # Bit totals for new lines for warped drawings need to be consistent with the
+            # current bit filler
+            bit_totals = []
             for b in range(n_bits):
-                if cur_layer_mid_measure is not None:
-                    # This condition is needed to pad the filler on the bits under MidMeasureMPs
-                    # that are used for conditions correctly
-                    cur_b_filler = (
-                        b_filler
-                        if bit_map[cur_layer_mid_measure] >= b and i < cwire_layers[b][-1]
-                        else " "
-                    )
-                else:
-                    cur_b_filler = b_filler if cwire_layers[b][0] < i < cwire_layers[b][-1] else " "
-                layer_str[b + n_wires] = layer_str[b + n_wires].ljust(max_label_len, cur_b_filler)
-
-            line_length += max_label_len + 1  # one for the filler character
-
-            if line_length > max_length:
-                # move totals into finished_lines and reset totals
-                finished_lines += wire_totals + bit_totals
-                finished_lines[-1] += "\n"
-                wire_totals = [w_filler] * n_wires
+                cur_b_filler = b_filler if cwire_layers[b][0] < i <= cwire_layers[b][-1] else " "
+                bit_totals.append(cur_b_filler)
 
-                # Bit totals for new lines for warped drawings need to be consistent with the
-                # current bit filler
-                bit_totals = []
-                for b in range(n_bits):
-                    cur_b_filler = (
-                        b_filler if cwire_layers[b][0] < i <= cwire_layers[b][-1] else " "
-                    )
-                    bit_totals.append(cur_b_filler)
+            line_length = 2 + max_label_len
 
-                line_length = 2 + max_label_len
+        ###################################################
+        # Join current layer with lines for previous layers
+        ###################################################
+        # Joining is done by adding a filler at the end of the previous layer
+        wire_totals = [w_filler.join([t, s]) for t, s in zip(wire_totals, layer_str[:n_wires])]
 
-            # Join current layer with lines for previous layers. Joining is done by adding a filler at
-            # the end of the previous layer
-            wire_totals = [w_filler.join([t, s]) for t, s in zip(wire_totals, layer_str[:n_wires])]
+        for j, (bt, s) in enumerate(zip(bit_totals, layer_str[n_wires : n_wires + n_bits])):
+            cur_b_filler = b_filler if cwire_layers[j][0] < i <= cwire_layers[j][-1] else " "
+            bit_totals[j] = cur_b_filler.join([bt, s])
 
-            for j, (bt, s) in enumerate(zip(bit_totals, layer_str[n_wires : n_wires + n_bits])):
-                cur_b_filler = b_filler if cwire_layers[j][0] < i <= cwire_layers[j][-1] else " "
-                bit_totals[j] = cur_b_filler.join([bt, s])
-
-        if ender:
+        ################################################
+        # Add ender characters to final operations layer
+        ################################################
+        if i == final_operations_layer:
             wire_totals = [f"{s}─┤" for s in wire_totals]
-            bit_totals = [f"{s}  " for s in bit_totals]
+            for b in range(n_bits):
+                if cwire_layers[b][-1] > final_operations_layer:
+                    bit_totals[b] += "═╡"
+                else:
+                    bit_totals[b] += "  "
 
             line_length += 2
 
diff --git a/pennylane/drawer/utils.py b/pennylane/drawer/utils.py
index 27a6d3399..9d474b94f 100644
--- a/pennylane/drawer/utils.py
+++ b/pennylane/drawer/utils.py
@@ -15,7 +15,7 @@
 This module contains some useful utility functions for circuit drawing.
 """
 from pennylane.ops import Controlled, Conditional
-from pennylane.measurements import MidMeasureMP
+from pennylane.measurements import MeasurementProcess, MidMeasureMP, MeasurementValue
 
 
 def default_wire_map(tape):
@@ -33,6 +33,40 @@ def default_wire_map(tape):
     return {wire: ind for ind, wire in enumerate(used_wires)}
 
 
+def default_bit_map(tape):
+    """Create a dictionary mapping ``MidMeasureMP``'s to indices corresponding to classical
+    wires. We only add mid-circuit measurements that are used for classical conditions and for
+    collecting statistics to this dictionary.
+
+    Args:
+        tape [~.tape.QuantumTape]: the QuantumTape containing operations and measurements
+
+    Returns:
+        dict: map from mid-circuit measurements to classical wires."""
+    bit_map = {}
+
+    for op in tape:
+        if isinstance(op, Conditional):
+            for m in op.meas_val.measurements:
+                bit_map[m] = None
+
+        if isinstance(op, MeasurementProcess) and op.mv is not None:
+            if isinstance(op.mv, MeasurementValue):
+                for m in op.mv.measurements:
+                    bit_map[m] = None
+            else:
+                for m in op.mv:
+                    bit_map[m.measurements[0]] = None
+
+    cur_cwire = 0
+    for op in tape:
+        if isinstance(op, MidMeasureMP) and op in bit_map:
+            bit_map[op] = cur_cwire
+            cur_cwire += 1
+
+    return bit_map
+
+
 def convert_wire_order(tape, wire_order=None, show_all_wires=False):
     """Creates the mapping between wire labels and place in order.
 
@@ -103,16 +137,18 @@ def unwrap_controls(op):
     return control_wires, control_values
 
 
-def cwire_connections(layers):
+def cwire_connections(layers, bit_map):
     """Extract the information required for classical control wires.
 
     Args:
         layers (List[List[.Operator, .MeasurementProcess]]): the operations and measurements sorted
             into layers via ``drawable_layers``. Measurement layers may be appended to operation layers.
+        bit_map (Dict): Dictionary containing mid-circuit measurements that are used for
+            classical conditions or measurement statistics as keys.
 
     Returns:
-        dict, list, list: map from mid circuit measurement to classical wire, list of list of accessed layers
-            for each classical wire, and largest wire corresponding to the accessed layers in the list above.
+        list, list: list of list of accessed layers for each classical wire, and largest wire
+        corresponding to the accessed layers in the list above.
 
     >>> with qml.queuing.AnnotatedQueue() as q:
     ...     m0 = qml.measure(0)
@@ -134,26 +170,32 @@ def cwire_connections(layers):
     layer will always be the one with the mid circuit measurement.
 
     """
-    bit_map = {}
-    for layer in layers:
-        for op in layer:
-            if isinstance(op, Conditional):
-                for m in op.meas_val.measurements:
-                    bit_map[m] = None  # place holder till next pass
+    if len(bit_map) == 0:
+        return [], []
 
     connected_layers = [[] for _ in bit_map]
     connected_wires = [[] for _ in bit_map]
-    num_cwires = 0
+
     for layer_idx, layer in enumerate(layers):
         for op in layer:
             if isinstance(op, MidMeasureMP) and op in bit_map:
-                bit_map[op] = num_cwires
-                connected_layers[num_cwires].append(layer_idx)
-                connected_wires[num_cwires].append(op.wires[0])
-                num_cwires += 1
+                connected_layers[bit_map[op]].append(layer_idx)
+                connected_wires[bit_map[op]].append(op.wires[0])
+
             elif isinstance(op, Conditional):
                 for m in op.meas_val.measurements:
                     cwire = bit_map[m]
                     connected_layers[cwire].append(layer_idx)
                     connected_wires[cwire].append(max(op.wires))
-    return bit_map, connected_layers, connected_wires
+
+            elif isinstance(op, MeasurementProcess) and op.mv is not None:
+                if isinstance(op.mv, MeasurementValue):
+                    for m in op.mv.measurements:
+                        cwire = bit_map[m]
+                        connected_layers[cwire].append(layer_idx)
+                else:
+                    for m in op.mv:
+                        cwire = bit_map[m.measurements[0]]
+                        connected_layers[cwire].append(layer_idx)
+
+    return connected_layers, connected_wires
diff --git a/pennylane/ops/functions/iterative_qpe.py b/pennylane/ops/functions/iterative_qpe.py
index b0136ccaa..6dcd5cb23 100644
--- a/pennylane/ops/functions/iterative_qpe.py
+++ b/pennylane/ops/functions/iterative_qpe.py
@@ -39,34 +39,39 @@ def iterative_qpe(base, ancilla, iters):
 
     .. code-block:: python
 
-        dev = qml.device("default.qubit", shots = 5)
+        dev = qml.device("default.qubit", shots=5)
 
         @qml.qnode(dev)
         def circuit():
 
           # Initial state
-          qml.PauliX(wires = [0])
+          qml.PauliX(wires=[0])
 
           # Iterative QPE
-          measurements = qml.iterative_qpe(qml.RZ(2., wires = [0]), ancilla = 1, iters = 3)
+          measurements = qml.iterative_qpe(qml.RZ(2.0, wires=[0]), ancilla=1, iters=3)
 
-          return [qml.sample(op = meas) for meas in measurements]
+          return qml.sample(measurements)
 
     .. code-block:: pycon
 
         >>> print(circuit())
-        [array([0, 0, 0, 0, 0]), array([1, 0, 0, 0, 0]), array([0, 1, 1, 1, 1])]
+        [[0 0 1]
+         [0 0 1]
+         [0 0 1]
+         [1 1 1]
+         [0 0 1]]
 
-    The output is an array of size ``(number of iterations, number of shots)``.
+    The output is an array of size ``(number of shots, number of iterations)``.
 
     .. code-block:: pycon
 
-        >>> print(qml.draw(circuit)())
+        >>> print(qml.draw(circuit, max_length=150)())
 
-        1: ──H─╭●────────────H──┤↗│  │0⟩──H─╭●────────────Rϕ(-1.57)──H──┤↗│  │0⟩──H─╭●────────────Rϕ(-1.57)──Rϕ(-0.79)──H──┤↗│  │0⟩─┤  Sample  Sample  Sample
-        0: ──X─╰RZ(2.00)⁴⋅⁰──────║──────────╰RZ(2.00)²⋅⁰──║──────────────║──────────╰RZ(2.00)¹⋅⁰──║──────────║──────────────────────┤
-                                 ╚════════════════════════╝══════════════║════════════════════════║══════════╝
-                                                                         ╚════════════════════════╝
+        0: ──X─╭RZ(2.00)⁴─────────────────╭RZ(2.00)²────────────────────────────╭RZ(2.00)¹────────────────────────────────────┤
+        1: ──H─╰●──────────H──┤↗│  │0⟩──H─╰●──────────Rϕ(-1.57)──H──┤↗│  │0⟩──H─╰●──────────Rϕ(-1.57)──Rϕ(-0.79)──H──┤↗│  │0⟩─┤
+                               ╚══════════════════════╩══════════════║══════════════════════║══════════╩══════════════║═══════╡ ╭Sample[MCM]
+                                                                     ╚══════════════════════╩═════════════════════════║═══════╡ ├Sample[MCM]
+                                                                                                                      ╚═══════╡ ╰Sample[MCM]
     """
 
     measurements = []


Test patch:
diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py
index d0da838f0..c9308ab78 100644
--- a/tests/drawer/test_draw.py
+++ b/tests/drawer/test_draw.py
@@ -576,10 +576,10 @@ class TestMidCircuitMeasurements:
 
         drawing = qml.draw(circ, max_length=25)()
         expected_drawing = (
-            "0: ──┤↗├─────────────\n"
+            "0: ──┤↗├──RX(0.00)───\n"
             "1: ───║───RX(0.12)──X\n"
             "      ╚═══╩═════════╝\n\n"
-            "───RX(0.00)─┤  <Z>\n"
+            "────────────┤  <Z>\n"
             "───RX(0.00)─┤     \n"
             "                  "
         )
@@ -668,6 +668,156 @@ class TestMidCircuitMeasurements:
 
         assert drawing == expected_drawing
 
+    @pytest.mark.parametrize(
+        "mp, mp_label",
+        [
+            (qml.expval, "<MCM>"),
+            (qml.var, "Var[MCM]"),
+            (qml.probs, "Probs[MCM]"),
+            (qml.sample, "Sample[MCM]"),
+            (qml.counts, "Counts[MCM]"),
+        ],
+    )
+    def test_single_meas_stats(self, mp, mp_label):
+        """Test that collecting statistics on a single mid-circuit measurement
+        works as expected."""
+
+        def circ():
+            qml.Hadamard(0)
+            m0 = qml.measure(0)
+            qml.Hadamard(0)
+            return mp(op=m0)
+
+        drawing = qml.draw(circ)()
+        expected_drawing = (
+            "0: ──H──┤↗├──H─┤  " + " " * len(mp_label) + f"\n         ╚═════╡  {mp_label}"
+        )
+
+        assert drawing == expected_drawing
+
+    @pytest.mark.parametrize(
+        "mp, mp_label",
+        [
+            (qml.expval, "<MCM>"),
+            (qml.var, "Var[MCM]"),
+            (qml.sample, "Sample[MCM]"),
+            (qml.counts, "Counts[MCM]"),
+        ],
+    )
+    def test_multi_meas_stats(self, mp, mp_label):
+        """Test that collecting statistics on multiple mid-circuit measurements
+        works as expected"""
+
+        def circ():
+            qml.Hadamard(0)
+            qml.Hadamard(1)
+            m0 = qml.measure(0)
+            m1 = qml.measure(1)
+            return mp(op=m0 + m1)
+
+        drawing = qml.draw(circ)()
+        expected_drawing = (
+            "0: ──H──┤↗├──────┤  "
+            + " " * len(mp_label)
+            + "\n1: ──H───║───┤↗├─┤  "
+            + " " * len(mp_label)
+            + f"\n         ╚════║══╡ ╭{mp_label}"
+            + f"\n              ╚══╡ ╰{mp_label}"
+        )
+
+        assert drawing == expected_drawing
+
+    def test_multi_meas_stats_multi_meas(self):
+        """Test that collecting statistics on multiple mid-circuit measurements with
+        multiple terminal measurement processes works as expected."""
+
+        def circ():
+            qml.Hadamard(0)
+            m0 = qml.measure(0)
+            qml.Hadamard(1)
+            m1 = qml.measure(1)
+            qml.Hadamard(2)
+            m2 = qml.measure(2)
+            return qml.expval(m0 * m2), qml.sample(m1)
+
+        drawing = qml.draw(circ)()
+        expected_drawing = (
+            "0: ──H──┤↗├─────────────────┤                    \n"
+            "1: ──────║───H──┤↗├─────────┤                    \n"
+            "2: ──────║───────║───H──┤↗├─┤                    \n"
+            "         ╚═══════║═══════║══╡ ╭<MCM>             \n"
+            "                 ╚═══════║══╡ │       Sample[MCM]\n"
+            "                         ╚══╡ ╰<MCM>             "
+        )
+
+        assert drawing == expected_drawing
+
+    def test_multi_meas_stats_same_cwire(self):
+        """Test that colecting statistics on multiple mid-circuit measurements
+        with multiple terminal measurements is drawn correctly"""
+
+        def circ():
+            qml.Hadamard(0)
+            qml.Hadamard(1)
+            m0 = qml.measure(0)
+            m1 = qml.measure(1)
+            return qml.expval(m0 + m1), qml.sample(m0)
+
+        drawing = qml.draw(circ)()
+        expected_drawing = (
+            "0: ──H──┤↗├──────┤                    \n"
+            "1: ──H───║───┤↗├─┤                    \n"
+            "         ╚════║══╡ ╭<MCM>  Sample[MCM]\n"
+            "              ╚══╡ ╰<MCM>             "
+        )
+
+        assert drawing == expected_drawing
+
+    def test_single_cond_single_meas_stats(self):
+        """Test that collecting statistics and using classical conditions together
+        on the same measurement works as expected."""
+
+        def circ():
+            qml.Hadamard(0)
+            m0 = qml.measure(0)
+            qml.cond(m0, qml.PauliZ)(0)
+            return qml.expval(m0)
+
+        drawing = qml.draw(circ)()
+        expected_drawing = "0: ──H──┤↗├──Z─┤       \n         ╚═══╩═╡  <MCM>"
+
+        assert drawing == expected_drawing
+
+    def test_multi_cond_multi_meas_stats(self):
+        """Test that combining multiple conditionals and multiple mid-circuit
+        measurement statistics is drawn correctly."""
+
+        def circ():
+            qml.Hadamard(0)
+            m0 = qml.measure(0)
+            qml.Hadamard(1)
+            m1 = qml.measure(1)
+            qml.Hadamard(2)
+            m2 = qml.measure(2)
+            qml.Hadamard(3)
+            qml.measure(3)
+            qml.cond(m0, qml.PauliX)(0)
+            qml.cond(m0 & m1, qml.PauliY)(1)
+            return qml.expval(m2), qml.sample([m1, m2])
+
+        drawing = qml.draw(circ)()
+        expected_drawing = (
+            "0: ──H──┤↗├──────────────────────────X────┤                    \n"
+            "1: ──────║───H──┤↗├──────────────────║──Y─┤                    \n"
+            "2: ──────║───────║───H──┤↗├──────────║──║─┤                    \n"
+            "3: ──────║───────║───────║───H──┤↗├──║──║─┤                    \n"
+            "         ╚═══════║═══════║═══════════╩══╣                      \n"
+            "                 ╚═══════║══════════════╩═╡        ╭Sample[MCM]\n"
+            "                         ╚════════════════╡  <MCM> ╰Sample[MCM]"
+        )
+
+        assert drawing == expected_drawing
+
 
 @pytest.mark.parametrize(
     "transform",
diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py
index d6ec97bf4..d6c2d68ec 100644
--- a/tests/drawer/test_draw_mpl.py
+++ b/tests/drawer/test_draw_mpl.py
@@ -341,8 +341,8 @@ def test_draw_mpl_with_qfunc_warns_with_expansion_strategy():
         _ = qml.draw_mpl(qfunc, expansion_strategy="gradient")
 
 
-def test_mid_circuit_measurement_device_api(mocker):
-    """Test that a circuit containing mid-circuit measurements is transformed by the drawer
+def test_qnode_mid_circuit_measurement_not_deferred_device_api(mocker):
+    """Test that a circuit containing mid-circuit measurements is not transformed by the drawer
     to use deferred measurements if the device uses the new device API."""
 
     @qml.qnode(qml.device("default.qubit"))
@@ -355,7 +355,7 @@ def test_mid_circuit_measurement_device_api(mocker):
     spy = mocker.spy(qml.defer_measurements, "_transform")
 
     _ = draw_qnode()
-    spy.assert_called_once()
+    spy.assert_not_called()
 
 
 def test_qnode_transform_program(mocker):
diff --git a/tests/drawer/test_drawable_layers.py b/tests/drawer/test_drawable_layers.py
index 113f0130c..263ac23a6 100644
--- a/tests/drawer/test_drawable_layers.py
+++ b/tests/drawer/test_drawable_layers.py
@@ -22,12 +22,13 @@ from pennylane.measurements import MidMeasureMP
 from pennylane.queuing import AnnotatedQueue
 from pennylane.drawer.drawable_layers import (
     _recursive_find_layer,
+    _recursive_find_mcm_stats_layer,
     drawable_layers,
 )
 
 
 class TestRecursiveFindLayer:
-    """Tests for `_recursive_find_layer`"""
+    """Tests for ``_recursive_find_layer`` and ``_recursive_find_mcm_stats_layer``"""
 
     def test_first_layer(self):
         """Test operation remains in 0th layer if not blocked"""
@@ -57,6 +58,55 @@ class TestRecursiveFindLayer:
         )
         assert out == 2
 
+    def test_first_mcm_stats_layer(self):
+        """Test operation remains in 0th layer if not blocked"""
+        mp0 = qml.measurements.MidMeasureMP(0, id="foo")
+        mp1 = qml.measurements.MidMeasureMP(1, id="bar")
+        bit_map = {mp0: 0, mp1: 1}
+
+        out = _recursive_find_mcm_stats_layer(
+            layer_to_check=0, op_occupied_cwires={0}, used_cwires_per_layer=[{1}], bit_map=bit_map
+        )
+        assert out == 0
+
+    def test_blocked_mcm_stats_layer(self):
+        """Test operation moved to higher layer if blocked on 0th layer."""
+        mp = qml.measurements.MidMeasureMP(0)
+        bit_map = {mp: 0}
+
+        out = _recursive_find_mcm_stats_layer(
+            layer_to_check=0, op_occupied_cwires={0}, used_cwires_per_layer=[{0}], bit_map=bit_map
+        )
+        assert out == 1
+
+    def test_recursion_no_block_mcm_stats(self):
+        """Test recursion to zero if start in higher layer and not blocked"""
+        mp0 = qml.measurements.MidMeasureMP(0)
+        mp1 = qml.measurements.MidMeasureMP(1)
+        bit_map = {mp0: 0, mp1: 1}
+
+        out = _recursive_find_mcm_stats_layer(
+            layer_to_check=2,
+            op_occupied_cwires={0},
+            used_cwires_per_layer=[{1}, {1}, {1}],
+            bit_map=bit_map,
+        )
+        assert out == 0
+
+    def test_recursion_block_mcm_stats(self):
+        """Test blocked on layer 1 gives placement on layer 2"""
+        mp0 = qml.measurements.MidMeasureMP(0)
+        mp1 = qml.measurements.MidMeasureMP(1)
+        bit_map = {mp0: 0, mp1: 1}
+
+        out = _recursive_find_mcm_stats_layer(
+            layer_to_check=3,
+            op_occupied_cwires={0},
+            used_cwires_per_layer=[{1}, {0}, {1}, {1}],
+            bit_map=bit_map,
+        )
+        assert out == 2
+
 
 class TestDrawableLayers:
     """Tests for `drawable_layers`"""
@@ -140,6 +190,10 @@ class TestDrawableLayers:
 
     def test_mid_measure_custom_wires(self):
         """Test that custom wires do not break the drawing of mid-circuit measurements."""
+        mp0 = MidMeasureMP("A", id="foo")
+        mp1 = MidMeasureMP("a", id="bar")
+        m0 = qml.measurements.MeasurementValue([mp0], lambda v: v)
+        m1 = qml.measurements.MeasurementValue([mp1], lambda v: v)
 
         def teleport(state):
             qml.QubitStateVector(state, wires=["A"])
@@ -147,15 +201,16 @@ class TestDrawableLayers:
             qml.CNOT(wires=["a", "B"])
             qml.CNOT(wires=["A", "a"])
             qml.Hadamard(wires="A")
-            m0 = qml.measure("A")
-            m1 = qml.measure("a")
+            qml.apply(mp0)
+            qml.apply(mp1)
             qml.cond(m1, qml.PauliX)("B")
             qml.cond(m0, qml.PauliZ)("B")
 
         tape_custom = qml.tape.make_qscript(teleport)([0, 1])
         [tape_standard], _ = qml.map_wires(tape_custom, {"A": 0, "a": 1, "B": 2})
         ops = tape_standard.operations
-        layers = drawable_layers(ops)
+        bit_map = {MidMeasureMP(0, id="foo"): None, MidMeasureMP(1, id="bar"): None}
+        layers = drawable_layers(ops, bit_map=bit_map)
         assert layers == [ops[:2]] + [[op] for op in ops[2:]]
 
 
@@ -168,7 +223,9 @@ class TestMidMeasure:
             m0 = qml.measure(0)
             qml.cond(m0, qml.PauliX)(1)
 
-        assert drawable_layers(q.queue) == [[q.queue[0]], [q.queue[1]]]
+        bit_map = {q.queue[0]: None}
+
+        assert drawable_layers(q.queue, bit_map=bit_map) == [[q.queue[0]], [q.queue[1]]]
 
     def test_cannot_draw_multi_wire_MidMeasureMP(self):
         """Tests that MidMeasureMP is only supported with one wire."""
diff --git a/tests/drawer/test_drawer_utils.py b/tests/drawer/test_drawer_utils.py
index 369a5af9d..902b18e92 100644
--- a/tests/drawer/test_drawer_utils.py
+++ b/tests/drawer/test_drawer_utils.py
@@ -19,6 +19,7 @@ import pytest
 import pennylane as qml
 from pennylane.drawer.utils import (
     default_wire_map,
+    default_bit_map,
     convert_wire_order,
     unwrap_controls,
     cwire_connections,
@@ -52,6 +53,32 @@ class TestDefaultWireMap:
         assert wire_map == {"a": 0, "b": 1, "c": 2}
 
 
+class TestDefaultBitMap:
+    """Tests ``default_bit_map`` helper function."""
+
+    def test_empty(self):
+        """Test creating empty bit map"""
+        # pylint: disable=use-implicit-booleaness-not-comparison
+        bit_map = default_bit_map([])
+        assert bit_map == {}
+
+        bit_map = default_bit_map([qml.measurements.MidMeasureMP(0)])
+        assert bit_map == {}
+
+    def test_simple(self):
+        """Test that the bit_map contains only measurements that are used."""
+
+        m0 = qml.measure(0)
+        m1 = qml.measure(1)
+        m2 = qml.measure(2)
+        cond0 = qml.ops.Conditional(m0, qml.S(0))
+        mp0 = qml.expval(m1)
+
+        queue = [m0.measurements[0], m1.measurements[0], m2.measurements[0], cond0, mp0]
+        bit_map = default_bit_map(queue)
+        assert bit_map == {m0.measurements[0]: 0, m1.measurements[0]: 1}
+
+
 class TestConvertWireOrder:
     """Tests the ``convert_wire_order`` utility function."""
 
@@ -170,15 +197,13 @@ class TestCwireConnections:
 
     def test_null_circuit(self):
         """Test null behavior with an empty circuit."""
-        cmap, layers, wires = cwire_connections([[]])
-        assert cmap == {}
+        layers, wires = cwire_connections([[]], {})
         assert layers == []
         assert wires == []
 
     def test_single_measure(self):
         """Test a single meassurment that does not have a conditional."""
-        cmap, layers, wires = cwire_connections([qml.measure(0).measurements])
-        assert cmap == {}
+        layers, wires = cwire_connections([qml.measure(0).measurements], {})
         assert layers == []
         assert wires == []
 
@@ -187,8 +212,9 @@ class TestCwireConnections:
         m = qml.measure(0)
         cond = qml.ops.Conditional(m, qml.PauliX(0))
         layers = [m.measurements, [cond]]
-        cmap, clayers, wires = cwire_connections(layers)
-        assert cmap == {m.measurements[0]: 0}
+        bit_map = {m.measurements[0]: 0}
+
+        clayers, wires = cwire_connections(layers, bit_map)
         assert clayers == [[0, 1]]
         assert wires == [[0, 0]]
 
@@ -200,10 +226,10 @@ class TestCwireConnections:
 
         cond0 = qml.ops.Conditional(m0 + m1, qml.PauliX(1))
         cond1 = qml.ops.Conditional(m1, qml.PauliY(2))
+        bit_map = {m0.measurements[0]: 0, m1.measurements[0]: 1}
 
         layers = [m0.measurements, m1.measurements, [cond0], m2_nonused.measurements, [cond1]]
-        cmap, clayers, wires = cwire_connections(layers)
-        assert cmap == {m0.measurements[0]: 0, m1.measurements[0]: 1}
+        clayers, wires = cwire_connections(layers, bit_map)
         assert clayers == [[0, 2], [1, 2, 4]]
         assert wires == [[0, 1], [1, 1, 2]]
 
@@ -213,7 +239,18 @@ class TestCwireConnections:
         m0 = qml.measure(0)
         cond0 = qml.ops.Conditional(m0, qml.S(0))
         layers = [m0.measurements, [cond0], [qml.expval(qml.PauliX(0))]]
-        cmap, clayers, wires = cwire_connections(layers)
-        assert cmap == {m0.measurements[0]: 0}
+        bit_map = {m0.measurements[0]: 0}
+        clayers, wires = cwire_connections(layers, bit_map)
         assert clayers == [[0, 1]]
         assert wires == [[0, 0]]
+
+    def test_mid_measure_stats_layer(self):
+        """Test cwire_connections works if layers contain terminal measurements using measurement
+        values"""
+
+        m0 = qml.measure(0)
+        layers = [m0.measurements, [qml.expval(m0)]]
+        bit_map = {m0.measurements[0]: 0}
+        clayers, wires = cwire_connections(layers, bit_map)
+        assert clayers == [[0, 1]]
+        assert wires == [[0]]
diff --git a/tests/drawer/test_tape_text.py b/tests/drawer/test_tape_text.py
index 627e9a462..372ca1875 100644
--- a/tests/drawer/test_tape_text.py
+++ b/tests/drawer/test_tape_text.py
@@ -25,6 +25,8 @@ from pennylane.drawer.tape_text import (
     _add_grouping_symbols,
     _add_cond_grouping_symbols,
     _add_mid_measure_grouping_symbols,
+    _add_cwire_measurement_grouping_symbols,
+    _add_cwire_measurement,
     _add_measurement,
     _add_op,
     _Config,
@@ -34,16 +36,21 @@ from pennylane.tape import QuantumScript, QuantumTape
 default_wire_map = {0: 0, 1: 1, 2: 2, 3: 3}
 default_bit_map = {}
 
-default_mid_measure_1 = qml.measurements.MidMeasureMP(0, id="foo")
-default_mid_measure_2 = qml.measurements.MidMeasureMP(0, id="bar")
+default_mid_measure_1 = qml.measurements.MidMeasureMP(0, id="1")
+default_mid_measure_2 = qml.measurements.MidMeasureMP(0, id="2")
+default_mid_measure_3 = qml.measurements.MidMeasureMP(0, id="3")
 default_measurement_value_1 = qml.measurements.MeasurementValue(
     [default_mid_measure_1], lambda v: v
 )
 default_measurement_value_2 = qml.measurements.MeasurementValue(
     [default_mid_measure_2], lambda v: v
 )
+default_measurement_value_3 = qml.measurements.MeasurementValue(
+    [default_mid_measure_3], lambda v: v
+)
 cond_bit_map_1 = {default_mid_measure_1: 0}
 cond_bit_map_2 = {default_mid_measure_1: 0, default_mid_measure_2: 1}
+stats_bit_map_1 = {default_mid_measure_1: 0, default_mid_measure_2: 1, default_mid_measure_3: 2}
 
 
 def get_conditional_op(mv, true_fn, *args, **kwargs):
@@ -171,6 +178,62 @@ class TestHelperFunctions:  # pylint: disable=too-many-arguments
             ),
         )
 
+    @pytest.mark.parametrize(
+        "mps, bit_map, out",
+        [
+            ([default_mid_measure_1], stats_bit_map_1, [" "] * (4 + len(stats_bit_map_1))),
+            (
+                [default_mid_measure_1, default_mid_measure_3],
+                stats_bit_map_1,
+                [" "] * len(default_wire_map) + ["╭", "│", "╰"],
+            ),
+            (
+                [default_mid_measure_1, default_mid_measure_2, default_mid_measure_3],
+                stats_bit_map_1,
+                [" "] * len(default_wire_map) + ["╭", "├", "╰"],
+            ),
+        ],
+    )
+    def test_add_cwire_measurement_grouping_symbols(self, mps, bit_map, out):
+        """Test private _add_cwire_measurement_grouping_symbols renders as expected."""
+        layer_str = [" "] * (len(default_wire_map) + len(bit_map))
+
+        assert out == _add_cwire_measurement_grouping_symbols(
+            mps, layer_str, _Config(wire_map=default_wire_map, bit_map=bit_map)
+        )
+
+    @pytest.mark.parametrize(
+        "mp, bit_map, out",
+        [
+            (
+                qml.sample([default_measurement_value_1]),
+                stats_bit_map_1,
+                [" "] * len(default_wire_map) + [" Sample[MCM]", " ", " "],
+            ),
+            (
+                qml.expval(default_measurement_value_1 * default_measurement_value_3),
+                stats_bit_map_1,
+                [" "] * len(default_wire_map) + ["╭<MCM>", "│", "╰<MCM>"],
+            ),
+            (
+                qml.counts(
+                    default_measurement_value_1
+                    + default_measurement_value_2
+                    - default_measurement_value_3
+                ),
+                stats_bit_map_1,
+                [" "] * len(default_wire_map) + ["╭Counts[MCM]", "├Counts[MCM]", "╰Counts[MCM]"],
+            ),
+        ],
+    )
+    def test_add_cwire_measurement(self, mp, bit_map, out):
+        """Test private _add_cwire_measurement renders as expected."""
+        layer_str = [" "] * (len(default_wire_map) + len(bit_map))
+
+        assert out == _add_cwire_measurement(
+            mp, layer_str, _Config(wire_map=default_wire_map, bit_map=bit_map)
+        )
+
     @pytest.mark.parametrize(
         "op, out",
         [