Repository: PennyLaneAI/pennylane
Pull request #4803 (2023-11-07 23:15:26)

Support drawing `qml.cond`: Part 1
**Context:**
Adding support for `qml.cond` for mid-circuit measurement drawing with `qml.draw`. This PR adds support for using one mid-circuit measurement with multiple `qml.cond`s, where each `qml.cond` uses the **same** mid-circuit measurement. An example of the drawing is shown below:

```python
import pennylane as qml

def circuit():
    qml.RX(0.5, 0)
    qml.RX(0.5, 1)
    m0 = qml.measure(0, reset=True, postselect=1)
    qml.cond(m0, qml.MultiControlledX)(wires=[1, 2, 3, 0], control_values="110")

    qml.CNOT([3, 2])
    qml.cond(m0, qml.ctrl(qml.MultiRZ, control=[1, 2], control_values=[True, False]))(0.5, wires=[0, 3])

    return qml.expval(qml.PauliZ(0))

print(qml.draw(circuit)())
```
```
0: ──RX(0.50)──┤↗₁│  │0⟩─╭X────╭MultiRZ(0.50)─┤  <Z>
1: ──RX(0.50)───║────────├●────├●─────────────┤     
2: ─────────────║────────├●─╭X─├○─────────────┤     
3: ─────────────║────────╰○─╰●─╰MultiRZ(0.50)─┤     
                ╚═════════╩═════╝                   
```

**Note:** This PR contains work that is foundational for follow up PRs but isn't used for drawing the cases mentioned above. I'll add arbitrary testing to get full coverage, but the aforementioned cases will be the only ones considered to currently be complete.

**Description of the Change:**

`drawable_layers`:
* `drawable_layers` no longer stacks mid-circuit measurements used by a `qml.cond` in the same layer.
* `drawable layers` now considers all wires under a mid-circuit measurement used for conditioning to be used.
* `drawable_layers` now considers all wires under a `qml.cond` to be used.

`tape_text`:
* Added helper functions to `tape_text.py` for correctly adding labels for mid-circuit measurements and `qml.cond`s, as well as for correctly padding the width of the layers for unused wires.
* Added helper function to `tape_text.py` for getting information about connectivity between mid-circuit measurements and conditional operations to keep track of when to start and stop drawing classical bits.
* `tape_text()` now tracks a `bit_map`, which stores a mapping between mid-circuit measurements and classical bits that are drawn for them. Classical bits are drawn using double lines `═`, and are placed under the wires in the drawing. Each mid-circuit measurement used for conditioning uses a classical bit. Measurements not used for conditioning are not allocated a bit.
* `tape_text()` now pads the classical bits between each layer of the drawing along with the wires.
* This PR also addresses #4809. The changes pertaining to that are contained in `tape_text.py::_add_measurements()`.

Other changes:
* Added `label` to `qml.transforms.Conditional`.
* Adjusted `qml.equal` to consider postselection when comparing `MidMeasureMP`s. This was missed in #4604.

**Benefits:**
Drawing support for mid-circuit measurements is better :)

**Possible Drawbacks:**
Code in `tape_text.py` is less readable. After a discussion with @trbromley , we decided that the best course of action is to leave refactoring as technical debt for the end of the release cycle. Reviews pertaining to major refactors are still appreciated :)

**Related GitHub Issues:**
#4809 

mudit2812: [sc-49112]
mudit2812: [sc-49764]
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4803?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `49 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`677fbe0`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/677fbe00f70e8456dacd5b0c4e9a46923036eeca?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64% compared to head [(`4cddf70`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4803?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.50%.

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


```diff
@@            Coverage Diff             @@
##           master    #4803      +/-   ##
==========================================
- Coverage   99.64%   99.50%   -0.15%     
==========================================
  Files         381      381              
  Lines       34377    34220     -157     
==========================================
- Hits        34256    34049     -207     
- Misses        121      171      +50     
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4803?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/ops/functions/equal.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4803?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9mdW5jdGlvbnMvZXF1YWwucHk=) | `98.55% <100.00%> (ø)` | |
| [pennylane/transforms/condition.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4803?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY29uZGl0aW9uLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/drawer/drawable\_layers.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4803?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci9kcmF3YWJsZV9sYXllcnMucHk=) | `98.50% <96.77%> (-1.50%)` | :arrow_down: |
| [pennylane/drawer/tape\_text.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4803?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci90YXBlX3RleHQucHk=) | `74.05% <56.75%> (-25.95%)` | :arrow_down: |

... and [43 files with indirect coverage changes](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4803/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4803?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).

trbromley: ```
0: ──RX(0.50)──┤↗₁│  │0⟩─╭X────╭MultiRZ(0.50)─┤  <Z>
1: ──RX(0.50)───║────────├●────├●─────────────┤     
2: ─────────────║────────├●─╭X─├○─────────────┤     
3: ─────────────║────────╰○─╰●─╰MultiRZ(0.50)─┤     
                ╚═════════╩═════╝                   
```
Have you considered the above? Specifically, the bottom classical line has `╩` in the middle rather than `╝`.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e444885f8..3c1219e5d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -46,8 +46,9 @@
 
 <h4>Drawing and statistics for mid-circuit measurements 🎨</h4>
 
-* `qml.draw` now supports drawing mid-circuit measurements.
+* `qml.draw` now supports drawing mid-circuit measurements and conditional operators.
   [(#4775)](https://github.com/PennyLaneAI/pennylane/pull/4775)
+  [(#4803)](https://github.com/PennyLaneAI/pennylane/pull/4803)
 
 <h4>Catalyst is seamlessly integrated with PennyLane ⚗️</h4>
 
@@ -246,6 +247,10 @@
   integer power.
   [(#4827)](https://github.com/PennyLaneAI/pennylane/pull/4827)
 
+* Fixed text drawer to correctly label qinfo measurements, as well as `qml.classical_shadow`
+  `qml.shadow_expval`.
+  [(#4803)](https://github.com/PennyLaneAI/pennylane/pull/4803)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/drawer/drawable_layers.py b/pennylane/drawer/drawable_layers.py
index 700db4273..cc8047e81 100644
--- a/pennylane/drawer/drawable_layers.py
+++ b/pennylane/drawer/drawable_layers.py
@@ -15,6 +15,7 @@
 This module contains a helper function to sort operations into layers.
 """
 
+from pennylane.ops import Conditional
 from pennylane.measurements import MidMeasureMP
 from .utils import default_wire_map
 
@@ -49,6 +50,38 @@ 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):
+    """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:
+            min_wire = mapped_wire
+            max_wire = max(wire_map.values())
+            return set(range(min_wire, max_wire + 1))
+
+        return {mapped_wire}
+
+    if isinstance(op, Conditional):
+        mapped_wires = [wire_map[wire] for wire in op.then_op.wires]
+        min_wire = min(mapped_wires)
+        max_wire = max(wire_map.values())
+        return set(range(min_wire, max_wire + 1))
+
+    if len(op.wires) == 0:
+        # if no wires, then it acts on all wires
+        # for example, qml.state and qml.sample
+        mapped_wires = set(wire_map.values())
+        return mapped_wires
+
+    mapped_wires = {wire_map[wire] for wire in op.wires}
+    # get all integers from the minimum to the maximum
+    min_wire = min(mapped_wires)
+    max_wire = max(mapped_wires)
+
+    return set(range(min_wire, max_wire + 1))
+
+
 def drawable_layers(ops, wire_map=None):
     """Determine non-overlapping yet dense placement of operations into layers for drawing.
 
@@ -91,56 +124,24 @@ def drawable_layers(ops, wire_map=None):
     occupied_wires_per_layer = [set()]
     ops_in_layer = [[]]
     measured_wires = {}
-    measured_layers = {}
+
+    # 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)
 
     # loop over operations
     for op in ops:
-        is_mid_measure = is_conditional = False
         if isinstance(op, MidMeasureMP):
             if len(op.wires) > 1:
                 raise ValueError("Cannot draw mid-circuit measurements with more than one wire.")
 
-            is_mid_measure = True
-            mapped_wire = wire_map[op.wires[0]]
-            measured_wires[op.id] = mapped_wire
-            op_occupied_wires = {mapped_wire}
-
-        elif len(op.wires) == 0:
-            # if no wires, then it acts on all wires
-            # for example, qml.state and qml.sample
-            mapped_wires = set(wire_map.values())
-            op_occupied_wires = mapped_wires
-        else:
-            mapped_wires = {wire_map[wire] for wire in op.wires}
-            if op.__class__.__name__ == "Conditional":
-                is_conditional = True
-                mapped_wires.update(measured_wires[m.id] for m in op.meas_val.measurements)
-            # get all integers from the minimum to the maximum
-            min_wire = min(mapped_wires)
-            max_wire = max(mapped_wires)
-            op_occupied_wires = set(range(min_wire, max_wire + 1))
+            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 is_conditional:
-            m_layers = [measured_layers[m.id] for m in op.meas_val.measurements]
-            max_control_layer = max(m_layers)
-
-            for m, m_layer in zip(op.meas_val.measurements, m_layers):
-                # TODO: remove from measurement caches for qubit reuse
-                # Note that along with this, wire lines will need to be re-drawn if being reused
-                #
-                # del measured_layers[m_id]
-                # del measured_wires[m_id]
-
-                # measurements for a conditional should all be in the same layer
-                if m_layer < max_control_layer:
-                    for i, o in enumerate(ops_in_layer[m_layer]):
-                        if m.id == o.id:
-                            ops_in_layer[max_control_layer].append(o)  # add to cond. op. layer
-                            ops_in_layer[m_layer].pop(i)  # remove from original layer
-                            # no need to remove from previous occupied_wires_per_layer
-
         # see if need to add new layer
         if op_layer > max_layer:
             max_layer += 1
@@ -150,7 +151,5 @@ def drawable_layers(ops, wire_map=None):
         # add to op_layer
         ops_in_layer[op_layer].append(op)
         occupied_wires_per_layer[op_layer].update(op_occupied_wires)
-        if is_mid_measure:
-            measured_layers[op.id] = op_layer
 
     return list(filter(None, ops_in_layer[:-1])) + ops_in_layer[-1:]
diff --git a/pennylane/drawer/tape_text.py b/pennylane/drawer/tape_text.py
index 09398281e..d77b48618 100644
--- a/pennylane/drawer/tape_text.py
+++ b/pennylane/drawer/tape_text.py
@@ -14,18 +14,46 @@
 """
 This module contains logic for the text based circuit drawer through the ``tape_text`` function.
 """
+# TODO: Fix the latter two pylint warnings
+# pylint: disable=too-many-arguments, too-many-branches, too-many-statements
 
+from dataclasses import dataclass
+from typing import Optional
 import pennylane as qml
-from pennylane.measurements import Expectation, Probability, Sample, Variance, State
+from pennylane.measurements import Expectation, Probability, Sample, Variance, State, MidMeasureMP
 
 from .drawable_layers import drawable_layers
-from .utils import convert_wire_order, unwrap_controls
+from .utils import convert_wire_order, unwrap_controls, find_mid_measure_cond_connections
 
 
-def _add_grouping_symbols(op, layer_str, wire_map):
+@dataclass
+class _Config:
+    """Dataclass containing attributes needed for updating the strings to be drawn for each layer"""
+
+    wire_map: dict
+    """Map between wire labels and their place in order"""
+
+    bit_map: dict
+    """Map between mid-circuit measurements and their corresponding bit in order"""
+
+    cur_layer: Optional[int] = None
+    """Current layer index that is being updated"""
+
+    final_cond_layers: Optional[list] = None
+    """List mapping bits to the last layer where they are used"""
+
+    decimals: Optional[int] = None
+    """Specifies how to round the parameters of operators"""
+
+    cache: Optional[dict] = None
+    """dictionary that carries information between label calls in the same drawing"""
+
+
+def _add_grouping_symbols(op, layer_str, config):  # pylint: disable=unused-argument
     """Adds symbols indicating the extent of a given object."""
+
     if len(op.wires) > 1:
-        mapped_wires = [wire_map[w] for w in op.wires]
+        mapped_wires = [config.wire_map[w] for w in op.wires]
         min_w, max_w = min(mapped_wires), max(mapped_wires)
         layer_str[min_w] = "╭"
         layer_str[max_w] = "╰"
@@ -36,27 +64,93 @@ def _add_grouping_symbols(op, layer_str, wire_map):
     return layer_str
 
 
-def _add_op(op, layer_str, wire_map, decimals, cache):
+def _add_cond_grouping_symbols(op, layer_str, config):
+    """Adds symbols indicating the extent of a given object for conditional
+    operators"""
+    n_wires = len(config.wire_map)
+
+    mapped_wires = [config.wire_map[w] for w in op.wires]
+    mapped_bits = [config.bit_map[m] for m in op.meas_val.measurements]
+    max_w = max(mapped_wires)
+    max_b = max(mapped_bits) + n_wires
+
+    ctrl_symbol = "╩" if config.cur_layer != config.final_cond_layers[max(mapped_bits)] else "╝"
+    layer_str[max_b] = "═" + ctrl_symbol
+
+    for w in range(max_w + 1, max(config.wire_map.values()) + 1):
+        layer_str[w] = "─║"
+
+    for b in range(n_wires, max_b):
+        if b - n_wires in mapped_bits:
+            layer_str[b] = "═╣"
+        else:
+            filler = " " if layer_str[b][-1] == " " else "═"
+            layer_str[b] = filler + "║"
+
+    return layer_str
+
+
+def _add_mid_measure_grouping_symbols(op, layer_str, config):
+    """Adds symbols indicating the extent of a given object for mid-measure
+    operators"""
+    if op not in config.bit_map:
+        return layer_str
+
+    n_wires = len(config.wire_map)
+    mapped_wire = config.wire_map[op.wires[0]]
+    bit = config.bit_map[op] + n_wires
+    layer_str[bit] += " ╚"
+
+    for w in range(mapped_wire + 1, n_wires):
+        layer_str[w] += "─║"
+
+    for b in range(n_wires, bit):
+        filler = " " if layer_str[b][-1] == " " else "═"
+        layer_str[b] += filler + "║"
+
+    return layer_str
+
+
+def _add_op(op, layer_str, config):
     """Updates ``layer_str`` with ``op`` operation."""
-    layer_str = _add_grouping_symbols(op, layer_str, wire_map)
+    if isinstance(op, qml.ops.Conditional):  # pylint: disable=no-member
+        layer_str = _add_cond_grouping_symbols(op, layer_str, config)
+        return _add_op(op.then_op, layer_str, config)
+
+    if isinstance(op, MidMeasureMP):
+        return _add_mid_measure_op(op, layer_str, config)
+
+    layer_str = _add_grouping_symbols(op, layer_str, config)
 
     control_wires, control_values = unwrap_controls(op)
 
     if control_values:
         for w, val in zip(control_wires, control_values):
-            layer_str[wire_map[w]] += "●" if val else "○"
+            layer_str[config.wire_map[w]] += "●" if val else "○"
     else:
         for w in control_wires:
-            layer_str[wire_map[w]] += "●"
+            layer_str[config.wire_map[w]] += "●"
 
-    label = op.label(decimals=decimals, cache=cache).replace("\n", "")
+    label = op.label(decimals=config.decimals, cache=config.cache).replace("\n", "")
     if len(op.wires) == 0:  # operation (e.g. barrier, snapshot) across all wires
         for i, s in enumerate(layer_str):
             layer_str[i] = s + label
     else:
         for w in op.wires:
             if w not in control_wires:
-                layer_str[wire_map[w]] += label
+                layer_str[config.wire_map[w]] += label
+
+    return layer_str
+
+
+def _add_mid_measure_op(op, layer_str, config):
+    """Updates ``layer_str`` with ``op`` operation when ``op`` is a
+    ``qml.measurements.MidMeasureMP``."""
+    layer_str = _add_mid_measure_grouping_symbols(op, layer_str, config)
+    label = op.label(decimals=config.decimals, cache=config.cache).replace("\n", "")
+
+    for w in op.wires:
+        layer_str[config.wire_map[w]] += label
 
     return layer_str
 
@@ -70,22 +164,25 @@ measurement_label_map = {
 }
 
 
-def _add_measurement(m, layer_str, wire_map, decimals, cache):
+def _add_measurement(m, layer_str, config):
     """Updates ``layer_str`` with the ``m`` measurement."""
-    layer_str = _add_grouping_symbols(m, layer_str, wire_map)
+    layer_str = _add_grouping_symbols(m, layer_str, config)
 
     if m.obs is None:
         obs_label = None
     else:
-        obs_label = m.obs.label(decimals=decimals, cache=cache).replace("\n", "")
-    meas_label = measurement_label_map[m.return_type](obs_label)
+        obs_label = m.obs.label(decimals=config.decimals, cache=config.cache).replace("\n", "")
+    if m.return_type in measurement_label_map:
+        meas_label = measurement_label_map[m.return_type](obs_label)
+    else:
+        meas_label = m.return_type.value
 
     if len(m.wires) == 0:  # state or probability across all wires
         for i, s in enumerate(layer_str):
             layer_str[i] = s + meas_label
 
     for w in m.wires:
-        layer_str[wire_map[w]] += meas_label
+        layer_str[config.wire_map[w]] += meas_label
     return layer_str
 
 
@@ -289,10 +386,6 @@ def tape_text(
     if n_wires == 0:
         return ""
 
-    totals = [f"{wire}: " for wire in wire_map]
-    line_length = max(len(s) for s in totals)
-    totals = [s.rjust(line_length, " ") for s in totals]
-
     # Used to store lines that are hitting the maximum length
     finished_lines = []
 
@@ -301,41 +394,115 @@ def tape_text(
         drawable_layers(tape.measurements, wire_map=wire_map),
     ]
     add_list = [_add_op, _add_measurement]
-    fillers = ["─", " "]
+    wire_fillers = ["─", " "]
+    bit_fillers = ["═", " "]
     enders = [True, False]  # add "─┤" after all operations
+    bit_maps = [{}, {}]
+
+    bit_maps[0], measurement_layers, final_cond_layers = find_mid_measure_cond_connections(
+        tape.operations, layers_list[0]
+    )
+    n_bits = len(bit_maps[0])
+
+    wire_totals = [f"{wire}: " for wire in wire_map]
+    bit_totals = ["" for _ in range(n_bits)]
+    line_length = max(len(s) for s in wire_totals)
+
+    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, bit_map in zip(
+        layers_list, add_list, wire_fillers, bit_fillers, enders, bit_maps
+    ):
+        # Collect information needed for drawing layers
+        config = _Config(
+            wire_map=wire_map,
+            bit_map=bit_map,
+            cur_layer=-1,
+            final_cond_layers=final_cond_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 measurement_layers[b] < i < final_cond_layers[b] else " "
+                layer_str[b + n_wires] = cur_b_filler
 
-    for layers, add, filler, ender in zip(layers_list, add_list, fillers, enders):
-        for layer in layers:
-            layer_str = [filler] * n_wires
+            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, wire_map)
+                    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, wire_map, decimals, cache)
+                    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)
-            layer_str = [s.ljust(max_label_len, filler) 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):
+                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 else " "
+                else:
+                    cur_b_filler = (
+                        b_filler if measurement_layers[b] < i < final_cond_layers[b] 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 += totals
+                finished_lines += wire_totals + bit_totals
                 finished_lines[-1] += "\n"
-                totals = [filler] * n_wires
+                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):
+                    cur_b_filler = (
+                        b_filler if measurement_layers[b] < i <= final_cond_layers[b] else " "
+                    )
+                    bit_totals.append(cur_b_filler)
+
                 line_length = 2 + max_label_len
 
-            totals = [filler.join([t, s]) for t, s in zip(totals, layer_str)]
+            # 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 measurement_layers[j] < i <= final_cond_layers[j] else " "
+                )
+                bit_totals[j] = cur_b_filler.join([bt, s])
+
         if ender:
-            totals = [s + "─┤" for s in totals]
+            wire_totals = [s + "─┤" for s in wire_totals]
+            bit_totals = [s + "  " for s in bit_totals]
+
             line_length += 2
 
     # Recursively handle nested tapes #
-    tape_totals = "\n".join(finished_lines + totals)
+    tape_totals = "\n".join(finished_lines + wire_totals + bit_totals)
     current_tape_offset = cache["tape_offset"]
     cache["tape_offset"] += len(tape_cache)
     for i, nested_tape in enumerate(tape_cache):
diff --git a/pennylane/drawer/utils.py b/pennylane/drawer/utils.py
index 603653eb8..b9a22c4ac 100644
--- a/pennylane/drawer/utils.py
+++ b/pennylane/drawer/utils.py
@@ -14,7 +14,8 @@
 """
 This module contains some useful utility functions for circuit drawing.
 """
-from pennylane.ops import Controlled
+from pennylane.ops import Controlled, Conditional
+from pennylane.measurements import MidMeasureMP
 
 
 def default_wire_map(tape):
@@ -100,3 +101,72 @@ def unwrap_controls(op):
 
     control_values = [bool(int(i)) for i in control_values] if control_values else control_values
     return control_wires, control_values
+
+
+def find_mid_measure_cond_connections(operations, layers):
+    """Collect and return information about connections between mid-circuit measurements
+    and classical conditions.
+
+    This utility function returns three items needed for processing mid-circuit measurements
+    and classical conditions for drawing:
+
+    * A dictionary mapping each mid-circuit measurement to a corresponding bit index.
+        This map only contains mid-circuit measurements that are used for classical conditioning.
+    * A list where each index is a bit and the values are the indices  of the layers containing
+        the mid-circuit measurement corresponding to the bits.
+    * A list where each index is a bit and the values are the indices of the last layers that
+        use those bits for classical conditions.
+
+    Args:
+        operations (list[~.Operation]): List of operations on the tape
+        layers (list[list[~.Operation]]): List of drawable layers containing list of operations
+            for each layer
+
+    Returns:
+        tuple[dict, list, list]: Data structures needed for correctly drawing classical conditions
+        as described above.
+    """
+
+    # Map between mid-circuit measurements and their position on the drawing
+    # The bit map only contains mid-circuit measurements that are used in
+    # classical conditions.
+    bit_map = {}
+
+    # Map between classical bit positions and the layer of their corresponding mid-circuit
+    # measurements.
+    measurement_layers = []
+
+    # Map between classical bit positions and the final layer where the bit is used.
+    # This is needed to know when to stop drawing a bit line. The bit is the index,
+    # so each of the two lists must have the same length as the number of bits
+    final_cond_layers = []
+
+    measurements_for_conds = set()
+    conditional_ops = []
+    for op in operations:
+        if isinstance(op, Conditional):
+            measurements_for_conds.update(op.meas_val.measurements)
+            conditional_ops.append(op)
+
+    if len(measurements_for_conds) > 0:
+        cond_mid_measures = [op for op in operations if op in measurements_for_conds]
+        cond_mid_measures.sort(key=operations.index)
+
+        bit_map = dict(zip(cond_mid_measures, range(len(cond_mid_measures))))
+
+        n_bits = len(bit_map)
+
+        # Set lists to correct size
+        measurement_layers = [None] * n_bits
+        final_cond_layers = [None] * n_bits
+
+        for i, layer in enumerate(layers):
+            for op in layer:
+                if isinstance(op, MidMeasureMP) and op in bit_map:
+                    measurement_layers[bit_map[op]] = i
+
+                if isinstance(op, Conditional):
+                    for mid_measure in op.meas_val.measurements:
+                        final_cond_layers[bit_map[mid_measure]] = i
+
+    return bit_map, measurement_layers, final_cond_layers
diff --git a/pennylane/ops/functions/equal.py b/pennylane/ops/functions/equal.py
index ed1a8fc62..df1fde2cc 100644
--- a/pennylane/ops/functions/equal.py
+++ b/pennylane/ops/functions/equal.py
@@ -384,7 +384,12 @@ def _equal_mid_measure(
     rtol=1e-5,
     atol=1e-9,
 ):
-    return op1.wires == op2.wires and op1.id == op2.id and op1.reset == op2.reset
+    return (
+        op1.wires == op2.wires
+        and op1.id == op2.id
+        and op1.reset == op2.reset
+        and op1.postselect == op2.postselect
+    )
 
 
 @_equal.register
diff --git a/pennylane/ops/op_math/condition.py b/pennylane/ops/op_math/condition.py
index e3da8c10a..a18dc7a15 100644
--- a/pennylane/ops/op_math/condition.py
+++ b/pennylane/ops/op_math/condition.py
@@ -52,6 +52,9 @@ class Conditional(Operation):
         self.then_op = then_op
         super().__init__(wires=then_op.wires, id=id)
 
+    def label(self, decimals=None, base_label=None, cache=None):
+        return self.then_op.label(decimals=decimals, base_label=base_label, cache=cache)
+
 
 def cond(condition, true_fn, false_fn=None):
     """Condition a quantum operation on the results of mid-circuit qubit measurements.


Test patch:
diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py
index 06cc0a646..0ffa8c291 100644
--- a/tests/drawer/test_draw.py
+++ b/tests/drawer/test_draw.py
@@ -281,77 +281,240 @@ class TestLayering:
         assert draw(circ)() == expected
 
 
-@pytest.mark.parametrize("device_name", ["default.qubit"])
-def test_mid_circuit_measurement_device_api(device_name, mocker):
-    """Test that a circuit containing mid-circuit measurements is transformed by the drawer
-    to use deferred measurements if the device uses the new device API."""
-    dev = qml.device(device_name)
+class TestMidCircuitMeasurements:
+    """Tests for drawing mid-circuit measurements and classical conditions."""
 
-    @qml.qnode(dev)
-    def circ():
-        qml.PauliX(0)
-        qml.measure(0)
-        return qml.probs(wires=0)
+    @pytest.mark.parametrize("device_name", ["default.qubit"])
+    def test_mid_circuit_measurement_device_api(self, device_name, mocker):
+        """Test that a circuit containing mid-circuit measurements is transformed by the drawer
+        to use deferred measurements if the device uses the new device API."""
+        dev = qml.device(device_name)
 
-    draw_qnode = qml.draw(circ)
-    spy = mocker.spy(qml.defer_measurements, "_transform")
+        @qml.qnode(dev)
+        def circ():
+            qml.PauliX(0)
+            qml.measure(0)
+            return qml.probs(wires=0)
+
+        draw_qnode = qml.draw(circ)
+        spy = mocker.spy(qml.defer_measurements, "_transform")
+
+        _ = draw_qnode()
+        spy.assert_called_once()
+
+    @pytest.mark.parametrize(
+        "postselect, reset, mid_measure_label",
+        [
+            (None, False, "┤↗├"),
+            (None, True, "┤↗│  │0⟩"),
+            (0, False, "┤↗₀├"),
+            (0, True, "┤↗₀│  │0⟩"),
+            (1, False, "┤↗₁├"),
+            (1, True, "┤↗₁│  │0⟩"),
+        ],
+    )
+    def test_draw_mid_circuit_measurement(self, postselect, reset, mid_measure_label):
+        """Test that mid-circuit measurements are drawn correctly."""
 
-    _ = draw_qnode()
-    spy.assert_called_once()
+        def func():
+            qml.Hadamard(0)
+            qml.measure(0, reset=reset, postselect=postselect)
+            qml.PauliX(0)
+            return qml.expval(qml.PauliZ(0))
 
+        drawing = qml.draw(func)()
+        expected_drawing = "0: ──H──" + mid_measure_label + "──X─┤  <Z>"
+
+        assert drawing == expected_drawing
+
+    def test_draw_mid_circuit_measurement_multiple_wires(self):
+        """Test that mid-circuit measurements are correctly drawn in circuits
+        with multiple wires."""
+
+        def circ(weights):
+            qml.RX(weights[0], 0)
+            qml.measure(0, reset=True)
+            qml.RX(weights[1], 1)
+            qml.measure(1)
+            qml.CNOT([0, 3])
+            qml.measure(3, postselect=0, reset=True)
+            qml.RY(weights[2], 2)
+            qml.CNOT([1, 2])
+            qml.measure(2, postselect=1)
+            qml.MultiRZ(0.5, [0, 2])
+            return qml.expval(qml.PauliZ(2))
+
+        drawing = qml.draw(circ)(np.array([np.pi, 3.124, 0.456]))
+        expected_drawing = (
+            "0: ──RX(3.14)──┤↗│  │0⟩─╭●─────────────────────╭MultiRZ(0.50)─┤     \n"
+            "1: ──RX(3.12)──┤↗├──────│─────────────╭●───────│──────────────┤     \n"
+            "3: ─────────────────────╰X──┤↗₀│  │0⟩─│────────│──────────────┤     \n"
+            "2: ──RY(0.46)─────────────────────────╰X──┤↗₁├─╰MultiRZ(0.50)─┤  <Z>"
+        )
 
-@pytest.mark.parametrize(
-    "postselect, reset, mid_measure_label",
-    [
-        (None, False, "┤↗├"),
-        (None, True, "┤↗│  │0⟩"),
-        (0, False, "┤↗₀├"),
-        (0, True, "┤↗₀│  │0⟩"),
-        (1, False, "┤↗₁├"),
-        (1, True, "┤↗₁│  │0⟩"),
-    ],
-)
-def test_draw_mid_circuit_measurement(postselect, reset, mid_measure_label):
-    """Test that mid-circuit measurements are drawn correctly."""
+        assert drawing == expected_drawing
 
-    def func():
-        qml.Hadamard(0)
-        qml.measure(0, reset=reset, postselect=postselect)
-        qml.PauliX(0)
-        return qml.expval(qml.PauliZ(0))
+    def test_single_meas_single_cond(self):
+        """Test that a circuit with a single classical condition and a single measurement
+        is drawn correctly."""
 
-    drawing = qml.draw(func)()
-    expected_drawing = "0: ──H──" + mid_measure_label + "──X─┤  <Z>"
-
-    assert drawing == expected_drawing
-
-
-def test_draw_mid_circuit_measurement_multiple_wires():
-    """Test that mid-circuit measurements are correctly drawn in circuits
-    with multiple wires."""
-
-    def circ(weights):
-        qml.RX(weights[0], 0)
-        qml.measure(0, reset=True)
-        qml.RX(weights[1], 1)
-        qml.measure(1)
-        qml.CNOT([0, 3])
-        qml.measure(3, postselect=0, reset=True)
-        qml.RY(weights[2], 2)
-        qml.CNOT([1, 2])
-        qml.measure(2, postselect=1)
-        qml.MultiRZ(0.5, [0, 2])
-        return qml.expval(qml.PauliZ(2))
-
-    drawing = qml.draw(circ)(np.array([np.pi, 3.124, 0.456]))
-    expected_drawing = (
-        "0: ──RX(3.14)──┤↗│  │0⟩─╭●─────────────────────╭MultiRZ(0.50)─┤     \n"
-        "1: ──RX(3.12)──┤↗├──────│─────────────╭●───────│──────────────┤     \n"
-        "3: ─────────────────────╰X──┤↗₀│  │0⟩─│────────│──────────────┤     \n"
-        "2: ──RY(0.46)─────────────────────────╰X──┤↗₁├─╰MultiRZ(0.50)─┤  <Z>"
-    )
+        def circ(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            qml.cond(m0, qml.PauliX)(wires=1)
+
+        drawing = qml.draw(circ)(np.pi)
+        expected_drawing = (
+            "0: ──RX(3.14)──┤↗├────┤  \n1: ─────────────║───X─┤  \n                ╚═══╝    "
+        )
+
+        assert drawing == expected_drawing
+
+    def test_single_meas_single_cond_multi_wire(self):
+        """Test that a multi-wire conditional operator relying on a single
+        mid-circuit measurement is drawn correctly."""
+
+        def circ(phi, theta):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            qml.RY(theta, 2)
+            qml.cond(m0, qml.CNOT)(wires=[1, 0])
+
+        drawing = qml.draw(circ)(np.pi, np.pi / 2)
+        expected_drawing = (
+            "0: ──RX(3.14)──┤↗├───────────╭X─┤  \n"
+            "2: ─────────────║───RY(1.57)─│──┤  \n"
+            "1: ─────────────║────────────╰●─┤  \n"
+            "                ╚═════════════╝    "
+        )
+
+        assert drawing == expected_drawing
+
+    def test_single_meas_multi_cond_with_reset(self):
+        """Test that a circuit where a single mid-circuit measurement is used for multiple
+        conditions is drawn correctly"""
+
+        def circ():
+            qml.RX(0.5, 0)
+            qml.RX(0.5, 1)
+            m0 = qml.measure(1, reset=True)
+            qml.cond(m0, qml.MultiControlledX)(wires=[1, 2, 0], control_values="10")
+            qml.CNOT([3, 2])
+            qml.cond(m0, qml.ctrl(qml.MultiRZ, control=[1, 2], control_values=[True, False]))(
+                0.5, wires=[0, 3]
+            )
+            qml.cond(m0, qml.PauliX)(0)
+            return qml.expval(qml.PauliZ(0))
+
+        drawing = qml.draw(circ)()
+        expected_drawing = (
+            "0: ──RX(0.50)───────────╭X────╭MultiRZ(0.50)──X─┤  <Z>\n"
+            "1: ──RX(0.50)──┤↗│  │0⟩─├●────├●──────────────║─┤     \n"
+            "2: ─────────────║───────╰○─╭X─├○──────────────║─┤     \n"
+            "3: ─────────────║────────║─╰●─╰MultiRZ(0.50)──║─┤     \n"
+            "                ╚════════╩═════╩══════════════╝       "
+        )
+
+        assert drawing == expected_drawing
+
+    def test_single_meas_multi_cond_with_postselection(self):
+        """Test that a circuit where a single mid-circuit measurement is used for multiple
+        conditions is drawn correctly"""
+
+        def circ():
+            qml.RX(0.5, 0)
+            qml.RX(0.5, 1)
+            m0 = qml.measure(1, postselect=1)
+            qml.cond(m0, qml.MultiControlledX)(wires=[1, 2, 0], control_values="10")
+            qml.CNOT([3, 2])
+            qml.cond(m0, qml.ctrl(qml.MultiRZ, control=[1, 2], control_values=[True, False]))(
+                0.5, wires=[0, 3]
+            )
+            qml.cond(m0, qml.PauliX)(0)
+            return qml.expval(qml.PauliZ(0))
+
+        drawing = qml.draw(circ)()
+        expected_drawing = (
+            "0: ──RX(0.50)───────╭X────╭MultiRZ(0.50)──X─┤  <Z>\n"
+            "1: ──RX(0.50)──┤↗₁├─├●────├●──────────────║─┤     \n"
+            "2: ─────────────║───╰○─╭X─├○──────────────║─┤     \n"
+            "3: ─────────────║────║─╰●─╰MultiRZ(0.50)──║─┤     \n"
+            "                ╚════╩═════╩══════════════╝       "
+        )
+
+        assert drawing == expected_drawing
+
+    def test_single_meas_multi_cond_split_lines(self):
+        """Test that a circuit is when multiple lines are needed and the measurement
+        and condition are on different lines."""
+
+        def circ():
+            m0 = qml.measure(0)
+            qml.RX(0.0, 0)
+            qml.cond(m0, qml.RX)(0.123, 1)
+            qml.RX(0.0, 0)
+            qml.cond(m0, qml.RX)(0.123, 1)
+            return qml.expval(qml.PauliZ(0))
+
+        drawing = qml.draw(circ, max_length=25)()
+        expected_drawing = (
+            "0: ──┤↗├──RX(0.00)\n"
+            "1: ───║───RX(0.12)\n"
+            "      ╚═══╩═══════\n\n"
+            "───RX(0.00)─┤  <Z>\n"
+            "───RX(0.12)─┤     \n"
+            "═══╝              "
+        )
+
+        assert drawing == expected_drawing
+
+    def test_single_meas_multi_cond_new_line(self):
+        """Test that a circuit is when multiple lines are needed and the measurement
+        and condition are on the same lines after the first line."""
+
+        def circ():
+            qml.RX(0.0, 0)
+            qml.RX(0.0, 0)
+            m0 = qml.measure(0)
+            qml.cond(m0, qml.PauliX)(1)
+            qml.cond(m0, qml.PauliX)(1)
+            return qml.expval(qml.PauliZ(0))
+
+        drawing = qml.draw(circ, max_length=25)()
+        expected_drawing = (
+            "0: ──RX(0.00)──RX(0.00)\n"
+            "1: ────────────────────\n"
+            "                       \n\n"
+            "───┤↗├───────┤  <Z>\n"
+            "────║───X──X─┤     \n"
+            "    ╚═══╩══╝       "
+        )
+
+        assert drawing == expected_drawing
+
+    def test_single_meas_multi_cond_first_line(self):
+        """Test that a circuit is when multiple lines are needed and the measurement
+        and condition are on the same lines in the first line."""
+
+        def circ():
+            m0 = qml.measure(0)
+            qml.cond(m0, qml.RX)(0.123, 1)
+            qml.cond(m0, qml.PauliX)(1)
+            qml.RX(0.0, 0)
+            qml.RX(0.0, 1)
+            return qml.expval(qml.PauliZ(0))
+
+        drawing = qml.draw(circ, max_length=25)()
+        expected_drawing = (
+            "0: ──┤↗├─────────────\n"
+            "1: ───║───RX(0.12)──X\n"
+            "      ╚═══╩═════════╝\n\n"
+            "───RX(0.00)─┤  <Z>\n"
+            "───RX(0.00)─┤     \n"
+            "                  "
+        )
 
-    assert drawing == expected_drawing
+        assert drawing == expected_drawing
 
 
 @pytest.mark.parametrize(
diff --git a/tests/drawer/test_drawable_layers.py b/tests/drawer/test_drawable_layers.py
index 5bac4268d..d56149074 100644
--- a/tests/drawer/test_drawable_layers.py
+++ b/tests/drawer/test_drawable_layers.py
@@ -150,29 +150,6 @@ class TestMidMeasure:
 
         assert drawable_layers(q.queue) == [[q.queue[0]], [q.queue[1]]]
 
-    def test_mid_measure_promoted_if_needed(self):
-        """Tests promotion of MidMeasureMPs if used for the same conditional."""
-        with AnnotatedQueue() as q:
-            m0 = qml.measure(0)
-            qml.Hadamard(1)
-            m1 = qml.measure(1)
-            qml.cond(m0 + m1, qml.PauliZ)(2)
-
-        ops = q.queue
-        assert drawable_layers(ops) == [[ops[1]], [ops[2], ops[0]], [ops[3]]]
-
-    def test_empty_layers_are_pruned(self):
-        """Tests that no empty layers are returned after MidMeasure promotion."""
-        with AnnotatedQueue() as q:
-            m0 = qml.measure(1)
-            qml.CNOT([0, 2])
-            m1 = qml.measure(0)
-            qml.cond(m0 + m1, qml.PauliZ)(2)
-
-        ops = q.queue
-        layers = drawable_layers(ops, wire_map={i: i for i in range(3)})
-        assert layers == [[ops[1]], [ops[2], ops[0]], [ops[3]]]
-
     def test_cannot_draw_multi_wire_MidMeasureMP(self):
         """Tests that MidMeasureMP is only supported with one wire."""
         with pytest.raises(ValueError, match="mid-circuit measurements with more than one wire."):
diff --git a/tests/drawer/test_drawer_utils.py b/tests/drawer/test_drawer_utils.py
index 6af59dca3..bdc13dc09 100644
--- a/tests/drawer/test_drawer_utils.py
+++ b/tests/drawer/test_drawer_utils.py
@@ -17,7 +17,12 @@ Unit tests for the pennylane.drawer.utils` module.
 
 import pytest
 import pennylane as qml
-from pennylane.drawer.utils import default_wire_map, convert_wire_order, unwrap_controls
+from pennylane.drawer.utils import (
+    default_wire_map,
+    convert_wire_order,
+    unwrap_controls,
+    find_mid_measure_cond_connections,
+)
 from pennylane.wires import Wires
 
 
@@ -157,3 +162,49 @@ class TestUnwrapControls:
 
         assert control_wires == expected_control_wires
         assert control_values == expected_control_values
+
+
+class TestFindMidMeasureCondConnections:
+    """Tests for the find_mid_measure_cond_connections helper function"""
+
+    def test_no_conds(self):
+        """Test that the return values are empty if there are no conditional operations."""
+        operations = [
+            qml.RX(0.123, 0),
+            qml.measurements.MidMeasureMP(0),
+            qml.RX(0.456, 1),
+            qml.measurements.MidMeasureMP(1),
+        ]
+        layers = [
+            [qml.RX(0.123, 0), qml.RX(0.456, 1)],
+            [qml.measurements.MidMeasureMP(0), qml.measurements.MidMeasureMP(1)],
+        ]
+
+        bit_map, measurement_layers, final_cond_layers = find_mid_measure_cond_connections(
+            operations, layers
+        )
+
+        assert bit_map == {}
+        assert len(measurement_layers) == len(final_cond_layers) == 0
+
+    def test_multi_meas_multi_cond(self):
+        """Test that multiple measurements and multiple conditions return the correct
+        output."""
+        with qml.queuing.AnnotatedQueue() as q:
+            m0 = qml.measure(0)
+            m1 = qml.measure(1)
+            m2 = qml.measure(2)
+            qml.cond(m0 & m1, qml.PauliX)(wires=4)
+            qml.cond(m1 & m2, qml.PauliX)(wires=5)
+            qml.cond(m2 & m0, qml.PauliX)(wires=6)
+
+        operations = q.queue
+        layers = [[op] for op in operations]
+
+        bit_map, measurement_layers, final_cond_layers = find_mid_measure_cond_connections(
+            operations, layers
+        )
+
+        assert bit_map == {operations[0]: 0, operations[1]: 1, operations[2]: 2}
+        assert measurement_layers == [0, 1, 2]
+        assert final_cond_layers == [5, 4, 5]
diff --git a/tests/drawer/test_tape_text.py b/tests/drawer/test_tape_text.py
index 2d32a2167..671300d12 100644
--- a/tests/drawer/test_tape_text.py
+++ b/tests/drawer/test_tape_text.py
@@ -21,10 +21,39 @@ import pytest
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.drawer import tape_text
-from pennylane.drawer.tape_text import _add_grouping_symbols, _add_measurement, _add_op
+from pennylane.drawer.tape_text import (
+    _add_grouping_symbols,
+    _add_cond_grouping_symbols,
+    _add_mid_measure_grouping_symbols,
+    _add_measurement,
+    _add_op,
+    _Config,
+)
 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_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
+)
+cond_bit_map_1 = {default_mid_measure_1: 0}
+cond_bit_map_2 = {default_mid_measure_1: 0, default_mid_measure_2: 1}
+
+
+def get_conditional_op(mv, true_fn, *args, **kwargs):
+    """Helper to get conditional operator."""
+
+    with qml.queuing.AnnotatedQueue() as q:
+        qml.cond(mv, true_fn)(*args, **kwargs)
+
+    return q.queue[0]
+
 
 with qml.queuing.AnnotatedQueue() as q_tape:
     qml.RX(1.23456, wires=0)
@@ -35,7 +64,7 @@ with qml.queuing.AnnotatedQueue() as q_tape:
 tape = qml.tape.QuantumScript.from_queue(q_tape)
 
 
-class TestHelperFunctions:
+class TestHelperFunctions:  # pylint: disable=too-many-arguments
     """Test helper functions for the tape text."""
 
     @pytest.mark.parametrize(
@@ -48,7 +77,83 @@ class TestHelperFunctions:
     )
     def test_add_grouping_symbols(self, op, out):
         """Test private _add_grouping_symbols function renders as expected."""
-        assert out == _add_grouping_symbols(op, ["", "", "", ""], default_wire_map)
+        assert out == _add_grouping_symbols(
+            op, ["", "", "", ""], _Config(wire_map=default_wire_map, bit_map=default_bit_map)
+        )
+
+    @pytest.mark.parametrize(
+        "op, bit_map, layer_str, out",
+        [
+            (default_mid_measure_1, default_bit_map, ["", "", "", ""], ["", "", "", ""]),
+            (
+                default_mid_measure_1,
+                cond_bit_map_1,
+                ["", "", "", "", ""],
+                ["", "─║", "─║", "─║", " ╚"],
+            ),
+            (
+                default_mid_measure_2,
+                cond_bit_map_2,
+                ["─", "─", "─", "─", " ", " "],
+                ["─", "──║", "──║", "──║", "  ║", "  ╚"],
+            ),
+        ],
+    )
+    def test_add_mid_measure_grouping_symbols(self, op, layer_str, bit_map, out):
+        """Test private _add_grouping_symbols function renders as expected for MidMeasureMPs."""
+        assert out == _add_mid_measure_grouping_symbols(
+            op, layer_str, _Config(wire_map=default_wire_map, bit_map=bit_map)
+        )
+
+    @pytest.mark.parametrize(
+        "cond_op, args, kwargs, out, bit_map, mv",
+        [
+            (
+                qml.PauliX,
+                [],
+                {"wires": 0},
+                ["─", "─║", "─║", "─║", "═╩"],
+                cond_bit_map_1,
+                default_measurement_value_1,
+            ),
+            (
+                qml.MultiRZ,
+                [0.5],
+                {"wires": [0, 1]},
+                ["─", "─", "─║", "─║", "═╩"],
+                cond_bit_map_1,
+                default_measurement_value_1,
+            ),
+            (
+                qml.Toffoli,
+                [],
+                {"wires": [0, 1, 2]},
+                ["─", "─", "─", "─║", " ║", "═╝"],
+                cond_bit_map_2,
+                default_measurement_value_2,
+            ),
+            (
+                qml.Toffoli,
+                [],
+                {"wires": [0, 1, 2]},
+                ["─", "─", "─", "─║", "═╣", "═╝"],
+                cond_bit_map_2,
+                default_measurement_value_1 & default_measurement_value_2,
+            ),
+        ],
+    )
+    def test_add_cond_grouping_symbols(self, cond_op, bit_map, mv, args, kwargs, out):
+        """Test private _add_grouping_symbols function renders as expected for Conditionals."""
+        op = get_conditional_op(mv, cond_op, *args, **kwargs)
+        layer_str = ["─", "─", "─", ""] + [" "] * len(bit_map)
+
+        assert out == _add_cond_grouping_symbols(
+            op,
+            layer_str,
+            _Config(
+                wire_map=default_wire_map, bit_map=bit_map, cur_layer=1, final_cond_layers=[0, 1]
+            ),
+        )
 
     @pytest.mark.parametrize(
         "op, out",
@@ -58,23 +163,38 @@ class TestHelperFunctions:
             (qml.var(qml.PauliX(1)), ["", "Var[X]", "", ""]),
             (qml.state(), ["State", "State", "State", "State"]),
             (qml.sample(), ["Sample", "Sample", "Sample", "Sample"]),
+            (qml.purity(0), ["purity", "", "", ""]),
+            (qml.vn_entropy([2, 1]), ["", "╭vnentropy", "╰vnentropy", ""]),
+            (qml.mutual_info(3, 1), ["", "╭mutualinfo", "│", "╰mutualinfo"]),
         ],
     )
     def test_add_measurements(self, op, out):
         """Test private _add_measurement function renders as expected."""
-        assert out == _add_measurement(op, [""] * 4, default_wire_map, None, None)
+        assert out == _add_measurement(
+            op, [""] * 4, _Config(wire_map=default_wire_map, bit_map=default_bit_map)
+        )
 
     def test_add_measurements_cache(self):
         """Test private _add_measurement function with a matrix cache."""
         cache = {"matrices": []}
         op = qml.expval(qml.Hermitian(np.eye(2), wires=0))
-        assert _add_measurement(op, ["", ""], {0: 0, 1: 1}, None, cache) == ["<𝓗(M0)>", ""]
+        assert _add_measurement(
+            op, ["", ""], _Config(wire_map={0: 0, 1: 1}, bit_map=default_bit_map, cache=cache)
+        ) == [
+            "<𝓗(M0)>",
+            "",
+        ]
 
         assert qml.math.allclose(cache["matrices"][0], np.eye(2))
 
         op2 = qml.expval(qml.Hermitian(np.eye(2), wires=1))
         # new op with same matrix, should have same M0 designation
-        assert _add_measurement(op2, ["", ""], {0: 0, 1: 1}, None, cache) == ["", "<𝓗(M0)>"]
+        assert _add_measurement(
+            op2, ["", ""], _Config(wire_map={0: 0, 1: 1}, bit_map=default_bit_map, cache=cache)
+        ) == [
+            "",
+            "<𝓗(M0)>",
+        ]
 
     @pytest.mark.parametrize(
         "op, out",
@@ -90,7 +210,73 @@ class TestHelperFunctions:
     )
     def test_add_op(self, op, out):
         """Test adding the first operation to array of strings"""
-        assert out == _add_op(op, ["─"] * 4, default_wire_map, None, None)
+        assert out == _add_op(
+            op, ["─"] * 4, _Config(wire_map=default_wire_map, bit_map=default_bit_map)
+        )
+
+    @pytest.mark.parametrize(
+        "op, bit_map, layer_str, out",
+        [
+            (default_mid_measure_1, default_bit_map, ["─", "─", "─", "─"], ["─┤↗├", "─", "─", "─"]),
+            (
+                default_mid_measure_1,
+                cond_bit_map_1,
+                ["─", "─", "─", "─", " "],
+                ["─┤↗├", "──║", "──║", "──║", "  ╚"],
+            ),
+            (
+                default_mid_measure_2,
+                cond_bit_map_2,
+                ["─", "─", "─", "─", " ", " "],
+                ["─┤↗├", "──║", "──║", "──║", "  ║", "  ╚"],
+            ),
+        ],
+    )
+    def test_add_mid_measure_op(self, op, layer_str, bit_map, out):
+        """Test adding the first MidMeasureMP to array of strings"""
+        assert out == _add_op(op, layer_str, _Config(wire_map=default_wire_map, bit_map=bit_map))
+
+    @pytest.mark.parametrize(
+        "cond_op, args, kwargs, out, bit_map, mv",
+        [
+            (
+                qml.MultiRZ,
+                [0.5],
+                {"wires": [0, 1]},
+                ["╭MultiRZ", "╰MultiRZ", "─║", "─║", "═╩"],
+                cond_bit_map_1,
+                default_measurement_value_1,
+            ),
+            (
+                qml.Toffoli,
+                [],
+                {"wires": [0, 1, 2]},
+                ["╭●", "├●", "╰X", "─║", "═╩"],
+                cond_bit_map_1,
+                default_measurement_value_1,
+            ),
+            (
+                qml.PauliX,
+                [],
+                {"wires": 1},
+                ["─", "─X", "─║", "─║", " ║", "═╝"],
+                cond_bit_map_2,
+                default_measurement_value_2,
+            ),
+        ],
+    )
+    def test_add_cond_op(self, cond_op, bit_map, mv, args, kwargs, out):
+        """Test adding the first Conditional to array of strings"""
+        op = get_conditional_op(mv, cond_op, *args, **kwargs)
+        layer_str = ["─", "─", "─", ""] + [" "] * len(bit_map)
+
+        assert out == _add_op(
+            op,
+            layer_str,
+            _Config(
+                wire_map=default_wire_map, bit_map=bit_map, cur_layer=1, final_cond_layers=[0, 1]
+            ),
+        )
 
     @pytest.mark.parametrize(
         "op, out",
@@ -102,18 +288,26 @@ class TestHelperFunctions:
     )
     def test_add_second_op(self, op, out):
         """Test adding a second operation to the array of strings"""
-        start = _add_op(qml.PauliX(0), ["─"] * 4, default_wire_map, None, None)
-        assert out == _add_op(op, start, default_wire_map, None, None)
+        start = _add_op(
+            qml.PauliX(0), ["─"] * 4, _Config(wire_map=default_wire_map, bit_map=default_bit_map)
+        )
+        assert out == _add_op(
+            op, start, _Config(wire_map=default_wire_map, bit_map=default_bit_map)
+        )
 
     def test_add_op_cache(self):
         """Test private _add_op method functions with a matrix cache."""
         cache = {"matrices": []}
         op1 = qml.QubitUnitary(np.eye(2), wires=0)
-        assert _add_op(op1, ["", ""], {0: 0, 1: 1}, None, cache) == ["U(M0)", ""]
+        assert _add_op(
+            op1, ["", ""], _Config(wire_map={0: 0, 1: 1}, bit_map=default_bit_map, cache=cache)
+        ) == ["U(M0)", ""]
 
         assert qml.math.allclose(cache["matrices"][0], np.eye(2))
         op2 = qml.QubitUnitary(np.eye(2), wires=1)
-        assert _add_op(op2, ["", ""], {0: 0, 1: 1}, None, cache) == ["", "U(M0)"]
+        assert _add_op(
+            op2, ["", ""], _Config(wire_map={0: 0, 1: 1}, bit_map=default_bit_map, cache=cache)
+        ) == ["", "U(M0)"]
 
 
 class TestEmptyTapes:
diff --git a/tests/ops/op_math/test_condition.py b/tests/ops/op_math/test_condition.py
index 689ba981c..020c51655 100644
--- a/tests/ops/op_math/test_condition.py
+++ b/tests/ops/op_math/test_condition.py
@@ -203,6 +203,21 @@ class TestCond:
             qml.cond(m_0, inp)()
 
 
+@pytest.mark.parametrize("op_class", [qml.PauliY, qml.Toffoli, qml.Hadamard, qml.CZ])
+def test_conditional_label(op_class):
+    """Test that the label for conditional oeprators is correct."""
+    base_op = op_class(wires=range(op_class.num_wires))
+
+    # Need to use queue because `qml.cond` doesn't return `Conditional` operators
+    with qml.queuing.AnnotatedQueue() as q:
+        m0 = qml.measure(0)
+        qml.cond(m0, op_class)(wires=range(op_class.num_wires))
+
+    cond_op = q.queue[1]
+
+    assert base_op.label() == cond_op.label()
+
+
 @pytest.mark.parametrize("terminal_measurement", terminal_meas)
 class TestOtherTransforms:
     """Tests that qml.cond works correctly with other transforms."""