Fixed drawer tests

pennylane/drawer/drawable_layers.py

@@ -16,6 +16,7 @@
 """
 
 from pennylane.measurements import MidMeasureMP
+from pennylane.tape import QuantumScript
 from .utils import default_wire_map
 
 
@@ -61,7 +62,7 @@ def _get_op_occupied_wires(op, wire_map, cond_measurements):
 
         return {mapped_wire}
 
-    if op.name == "Conditional":
+    if op.__class__.__name__ == "Conditional":
         mapped_wires = [wire_map[wire] for wire in op.then_op.wires]
         min_wire = min(mapped_wires)
         max_wire = max(wire_map.values())
@@ -129,7 +130,7 @@ def drawable_layers(ops, wire_map=None):
     # Collect all mid-circuit measurements used for classical conditioning
     cond_measurements = set()
     for op in ops:
-        if op.name == "Conditional":
+        if op.__class__.__name__ == "Conditional":
             cond_measurements.update(op.meas_val.measurements)
 
     # loop over operations
@@ -143,7 +144,7 @@ def drawable_layers(ops, wire_map=None):
             is_mid_measure = True
             measured_wires[op.id] = wire_map[op.wires[0]]
 
-        elif op.name == "Conditional":
+        elif op.__class__.__name__ == "Conditional":
             is_conditional = True
 
         op_occupied_wires = _get_op_occupied_wires(op, wire_map, cond_measurements)

pennylane/drawer/tape_text.py

@@ -19,6 +19,7 @@
 
 import pennylane as qml
 from pennylane.measurements import Expectation, Probability, Sample, Variance, State, MidMeasureMP
+from pennylane.tape import QuantumScript
 
 from .drawable_layers import drawable_layers
 from .utils import convert_wire_order, unwrap_controls
@@ -86,7 +87,7 @@ def _add_mid_measure_grouping_symbols(op, layer_str, wire_map, bit_map):
 
 def _add_op(op, layer_str, wire_map, bit_map, decimals, cache):
     """Updates ``layer_str`` with ``op`` operation."""
-    if op.name == "Conditional":
+    if op.__class__.__name__ == "Conditional":
         layer_str = _add_cond_grouping_symbols(op, layer_str, wire_map, bit_map)
         return _add_op(op.then_op, layer_str, wire_map, bit_map, decimals, cache)
 
@@ -181,7 +182,7 @@ def _find_mid_measure_cond_connections(operations, layers):
     measurements_for_conds = set()
     conditional_ops = []
     for op in operations:
-        if op.name == "Conditional":
+        if op.__class__.__name__ == "Conditional":
             measurements_for_conds.update(op.meas_val.measurements)
             conditional_ops.append(op)
 
@@ -206,7 +207,7 @@ def _find_mid_measure_cond_connections(operations, layers):
         # using those layers
         for i, layer in enumerate(layers[0]):
             for op in layer:
-                if op.name == "Conditional":
+                if op.__class__.__name__ == "Conditional":
                     for mid_measure in op.meas_val.measurements:
                         all_bit_terminal_layers[0][bit_map[mid_measure]] = i
 

tests/drawer/test_drawable_layers.py

@@ -150,29 +150,6 @@ def test_basic_mid_measure(self):
 
         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."):

tests/drawer/test_tape_text.py

@@ -25,6 +25,7 @@
 from pennylane.tape import QuantumScript, QuantumTape
 
 default_wire_map = {0: 0, 1: 1, 2: 2, 3: 3}
+default_bit_map = {}
 
 with qml.queuing.AnnotatedQueue() as q_tape:
     qml.RX(1.23456, wires=0)
@@ -48,7 +49,7 @@ 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, ["", "", "", ""], default_wire_map, default_bit_map)
 
     @pytest.mark.parametrize(
         "op, out",
@@ -62,19 +63,25 @@ def test_add_grouping_symbols(self, op, out):
     )
     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, default_wire_map, default_bit_map, None, None)
 
     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, ["", ""], {0: 0, 1: 1}, default_bit_map, None, 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, ["", ""], {0: 0, 1: 1}, default_bit_map, None, cache) == [
+            "",
+            "<𝓗(M0)>",
+        ]
 
     @pytest.mark.parametrize(
         "op, out",
@@ -90,7 +97,7 @@ def test_add_measurements_cache(self):
     )
     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, default_wire_map, default_bit_map, None, None)
 
     @pytest.mark.parametrize(
         "op, out",
@@ -102,18 +109,18 @@ def test_add_op(self, op, out):
     )
     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, default_wire_map, default_bit_map, None, None)
+        assert out == _add_op(op, start, default_wire_map, default_bit_map, None, None)
 
     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, ["", ""], {0: 0, 1: 1}, default_bit_map, None, 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, ["", ""], {0: 0, 1: 1}, default_bit_map, None, cache) == ["", "U(M0)"]
 
 
 class TestEmptyTapes: