Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] `draw_mpl` is not fully compatible with work wires
### Expected behavior

A circuit can be drawn even if the operators contain work wires (that are not to be drawn!)

### Actual behavior

Mapping around wires fails in `draw_mpl` if there are work wires present that get overwritten.


### Additional information

This is quite similar to #7580 , fixed in #7581 .

### Source code

```shell
import pennylane as qml


dev = qml.device("default.qubit")
@qml.qnode(dev)
def circuit():

    qml.ctrl(qml.CRZ(0.5, [0, 2]), control=[3], work_wires=[1])
    return qml.probs(wires=[3])

qml.draw_mpl(circuit)()
```

### Tracebacks

```shell
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[5], line 11
      8     qml.ctrl(qml.CRZ(0.5, [0, 2]), control=[3], work_wires=[1])
      9     return qml.probs(wires=[3])
---> 11 qml.draw_mpl(circuit)()

File ~/repos/pennylane/pennylane/drawer/draw.py:858, in _draw_mpl_qnode.<locals>.wrapper(*args, **kwargs_qnode)
    855     except TypeError:
    856         _wire_order = tape.wires
--> 858 return tape_mpl(
    859     tape,
    860     wire_order=_wire_order,
    861     show_all_wires=show_all_wires,
    862     decimals=decimals,
    863     style=style,
    864     fig=fig,
    865     **kwargs,
    866 )

File ~/repos/pennylane/pennylane/drawer/tape_mpl.py:545, in tape_mpl(tape, wire_order, show_all_wires, decimals, style, fig, max_length, **kwargs)
    542     _set_style(style)
    544 try:
--> 545     return _tape_mpl(
    546         tape,
    547         wire_order=wire_order,
    548         show_all_wires=show_all_wires,
    549         decimals=decimals,
    550         fig=fig,
    551         max_length=max_length,
    552         **kwargs,
    553     )
    554 finally:
    555     if update_style:
    556         # we don't want to mess with how it modifies whether the interface is interactive
    557         # but we want to restore everything else

File ~/repos/pennylane/pennylane/drawer/tape_mpl.py:294, in _tape_mpl(tape, wire_order, show_all_wires, max_length, **kwargs)
    292 wire_map = convert_wire_order(tape, wire_order=wire_order, show_all_wires=show_all_wires)
    293 tape = transform_deferred_measurements_tape(tape)
--> 294 tape = ops.functions.map_wires(tape, wire_map=wire_map)[0][0]
    295 bit_map = default_bit_map(tape)
    297 layers = drawable_layers(tape.operations, wire_map={i: i for i in tape.wires}, bit_map=bit_map)

File ~/repos/pennylane/pennylane/ops/functions/map_wires.py:229, in map_wires(input, wire_map, queue, replace)
    227         return new_op
    228     return input.map_wires(wire_map=wire_map)
--> 229 return _map_wires_transform(input, wire_map=wire_map, queue=queue)

File ~/repos/pennylane/pennylane/transforms/core/transform_dispatcher.py:254, in TransformDispatcher.__call__(self, *targs, **tkwargs)
    251         return expand_processing(processed_results)
    253 else:
--> 254     transformed_tapes, processing_fn = self._transform(obj, *targs, **tkwargs)
    256 if self.is_informative:
    257     return processing_fn(transformed_tapes)

File ~/repos/pennylane/pennylane/ops/functions/map_wires.py:241, in _map_wires_transform(tape, wire_map, queue)
    237 @partial(transform, plxpr_transform=map_wires_plxpr_to_plxpr)
    238 def _map_wires_transform(
    239     tape: QuantumScript, wire_map=None, queue=False
    240 ) -> tuple[QuantumScriptBatch, PostprocessingFn]:
--> 241     ops = [
    242         (
    243             map_wires(op, wire_map, queue=queue)
    244             if not isinstance(op, QuantumScript)
    245             else map_wires(op, wire_map, queue=queue)[0][0]
    246         )
    247         for op in tape.operations
    248     ]
    249     measurements = [map_wires(m, wire_map, queue=queue) for m in tape.measurements]
    251     out = tape.__class__(
    252         ops=ops, measurements=measurements, shots=tape.shots, trainable_params=tape.trainable_params
    253     )

File ~/repos/pennylane/pennylane/ops/functions/map_wires.py:243, in <listcomp>(.0)
    237 @partial(transform, plxpr_transform=map_wires_plxpr_to_plxpr)
    238 def _map_wires_transform(
    239     tape: QuantumScript, wire_map=None, queue=False
    240 ) -> tuple[QuantumScriptBatch, PostprocessingFn]:
    241     ops = [
    242         (
--> 243             map_wires(op, wire_map, queue=queue)
    244             if not isinstance(op, QuantumScript)
    245             else map_wires(op, wire_map, queue=queue)[0][0]
    246         )
    247         for op in tape.operations
    248     ]
    249     measurements = [map_wires(m, wire_map, queue=queue) for m in tape.measurements]
    251     out = tape.__class__(
    252         ops=ops, measurements=measurements, shots=tape.shots, trainable_params=tape.trainable_params
    253     )

File ~/repos/pennylane/pennylane/ops/functions/map_wires.py:228, in map_wires(input, wire_map, queue, replace)
    226             qml.apply(new_op)
    227         return new_op
--> 228     return input.map_wires(wire_map=wire_map)
    229 return _map_wires_transform(input, wire_map=wire_map, queue=queue)

File ~/repos/pennylane/pennylane/ops/op_math/controlled.py:699, in Controlled.map_wires(self, wire_map)
    696 new_control_wires = Wires([wire_map.get(wire, wire) for wire in self.control_wires])
    697 new_work_wires = Wires([wire_map.get(wire, wire) for wire in self.work_wires])
--> 699 return ctrl(
    700     op=new_base,
    701     control=new_control_wires,
    702     control_values=self.control_values,
    703     work_wires=new_work_wires,
    704 )

File ~/repos/pennylane/pennylane/ops/op_math/controlled.py:169, in ctrl(op, control, control_values, work_wires, work_wire_type)
    161 if math.is_abstract(op):
    162     return Controlled(
    163         op,
    164         control,
   (...)
    167         work_wire_type=work_wire_type,
    168     )
--> 169 return create_controlled_op(
    170     op,
    171     control,
    172     control_values=control_values,
    173     work_wires=work_wires,
    174     work_wire_type=work_wire_type,
    175 )

File ~/repos/pennylane/pennylane/ops/op_math/controlled.py:229, in create_controlled_op(op, control, control_values, work_wires, work_wire_type)
    220     return ctrl(
    221         op.base,
    222         control=control + op.control_wires,
   (...)
    225         work_wire_type=work_wire_type,
    226     )
    228 if isinstance(op, Operator):
--> 229     return Controlled(
    230         op,
    231         control_wires=control,
    232         control_values=control_values,
    233         work_wires=work_wires,
    234         work_wire_type=work_wire_type,
    235     )
    237 if not callable(op):
    238     raise ValueError(
    239         f"The object {op} of type {type(op)} is not an Operator or callable. "
    240         "This error might occur if you apply ctrl to a list "
    241         "of operations instead of a function or Operator."
    242     )

File ~/repos/pennylane/pennylane/capture/capture_meta.py:30, in _stop_autograph.<locals>.new_f(*args, **kwargs)
     29 def new_f(*args, **kwargs):
---> 30     return f(*args, **kwargs)

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

File ~/repos/pennylane/pennylane/ops/op_math/controlled.py:1048, in ControlledOp.__init__(self, base, control_wires, control_values, work_wires, work_wire_type, id)
   1039 def __init__(
   1040     self,
   1041     base,
   (...)
   1046     id=None,
   1047 ):
-> 1048     super().__init__(base, control_wires, control_values, work_wires, work_wire_type, id)
   1049     # check the grad_recipe validity
   1050     if self.grad_recipe is None:
   1051         # Make sure grad_recipe is an iterable of correct length instead of None

File ~/repos/pennylane/pennylane/ops/op_math/controlled.py:597, in Controlled.__init__(self, base, control_wires, control_values, work_wires, work_wire_type, id)
    594     raise ValueError("The control wires must be different from the base operation wires.")
    596 if len(Wires.shared_wires([work_wires, base.wires + control_wires])) != 0:
--> 597     raise ValueError(
    598         "Work wires must be different the control_wires and base operation wires."
    599     )
    601 if work_wire_type not in {"clean", "dirty"}:
    602     raise ValueError(
    603         f"work_wire_type must be either 'clean' or 'dirty'. Got '{work_wire_type}'."
    604     )

ValueError: Work wires must be different the control_wires and base operation wires.
```

### System information

```shell
pl dev
```

### Existing GitHub issues

- [x] I have searched existing GitHub issues to make sure the issue does not already exist.

Pull request description:
Fix `qml.draw_mpl` with `work_wires` present
**Context:**
`qml.draw_mpl` maps wires before drawing. 
It does not consider work wires, so that there can arise collisions between mapped algorithmic wires and non-mapped work wires. This is similar to the problem in #7580 (fixed in #7581).

**Description of the Change:**
Consider work wires in the wire mapping.
As the map is also used to decide which wires are drawn at all, we need to now produce two maps: One for mapping (now incl. work wires) and one for indicating which wires are drawn (excl. work wires).

**Benefits:**
fixed bug.

**Possible Drawbacks:**
n/a

**Related GitHub Issues:**
fixes #7659 .

[sc-93132]


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 6a4d6a6e5..cad043c9b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -806,6 +806,11 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 
 <h3>Bug fixes 🐛</h3>
 
+* A bug in `qml.draw_mpl` for circuits with work wires has been fixed. The previously
+  inconsistent mapping for these wires has been resolved, ensuring accurate assignment during
+  drawing.
+  [(#7668)](https://github.com/PennyLaneAI/pennylane/pull/7668)
+
 * A bug in `ops.op_math.Prod.simplify()` has been fixed that led to global phases being discarded
   in special cases. Concretely, this problem occurs when Pauli factors combine into the identity
   up to a global phase _and_ there is no Pauli representation of the product operator.
diff --git a/pennylane/drawer/drawable_layers.py b/pennylane/drawer/drawable_layers.py
index 98a916529..9a4b26758 100644
--- a/pennylane/drawer/drawable_layers.py
+++ b/pennylane/drawer/drawable_layers.py
@@ -149,7 +149,7 @@ def drawable_layers(operations, wire_map=None, bit_map=None):
 
     """
 
-    wire_map = wire_map or default_wire_map(operations)
+    wire_map = wire_map or default_wire_map(operations)[1]
     bit_map = bit_map or {}
 
     # initialize for operation layers
diff --git a/pennylane/drawer/tape_mpl.py b/pennylane/drawer/tape_mpl.py
index 1ef83cab0..cc1477b78 100644
--- a/pennylane/drawer/tape_mpl.py
+++ b/pennylane/drawer/tape_mpl.py
@@ -289,9 +289,11 @@ def _tape_mpl(tape, wire_order=None, show_all_wires=False, max_length=None, **kw
         When max_length is None: tuple(Figure, Axes) for the complete circuit
         Otherwise: list[tuple(Figure, Axes)] with one pair per circuit segment
     """
-    wire_map = convert_wire_order(tape, wire_order=wire_order, show_all_wires=show_all_wires)
+    full_wire_map, used_wire_map = convert_wire_order(
+        tape, wire_order=wire_order, show_all_wires=show_all_wires
+    )
     tape = transform_deferred_measurements_tape(tape)
-    tape = ops.functions.map_wires(tape, wire_map=wire_map)[0][0]
+    tape = ops.functions.map_wires(tape, wire_map=full_wire_map)[0][0]
     bit_map = default_bit_map(tape)
 
     layers = drawable_layers(tape.operations, wire_map={i: i for i in tape.wires}, bit_map=bit_map)
@@ -311,7 +313,9 @@ def _tape_mpl(tape, wire_order=None, show_all_wires=False, max_length=None, **kw
     )
 
     if max_length is None:
-        return _draw_layers(layers, tape.measurements, config=config, wire_map=wire_map, **kwargs)
+        return _draw_layers(
+            layers, tape.measurements, config=config, wire_map=used_wire_map, **kwargs
+        )
 
     layer_count = len(layers)
     return [
@@ -319,7 +323,7 @@ def _tape_mpl(tape, wire_order=None, show_all_wires=False, max_length=None, **kw
             layers=layers[i : i + max_length],
             measurements=tape.measurements if i + max_length >= layer_count else "dots",
             config=config,
-            wire_map=wire_map,
+            wire_map=used_wire_map,
             starting_dots=i > 0,
             **kwargs,
         )
diff --git a/pennylane/drawer/tape_text.py b/pennylane/drawer/tape_text.py
index 5db64c61b..ffb1d2c86 100644
--- a/pennylane/drawer/tape_text.py
+++ b/pennylane/drawer/tape_text.py
@@ -414,7 +414,7 @@ def tape_text(
     cache.setdefault("matrices", [])
     tape_cache = []
 
-    wire_map = convert_wire_order(tape, wire_order=wire_order, show_all_wires=show_all_wires)
+    _, 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)
     if n_wires == 0:
diff --git a/pennylane/drawer/utils.py b/pennylane/drawer/utils.py
index 8a33479d0..6d39796e1 100644
--- a/pennylane/drawer/utils.py
+++ b/pennylane/drawer/utils.py
@@ -27,12 +27,19 @@ def default_wire_map(tape):
         tape [~.tape.QuantumTape): the QuantumTape containing operations and measurements
 
     Returns:
-        dict: map from wires to sequential positive integers
+        tuple[dict]: A tuple of maps from wires to sequential positive integers. The first map
+        includes work wires whereas the second map excludes work wires.
     """
 
     # Use dictionary to preserve ordering, sets break order
     used_wires = {wire: None for op in tape for wire in op.wires}
-    return {wire: ind for ind, wire in enumerate(used_wires)}
+    used_wire_map = {wire: ind for ind, wire in enumerate(used_wires)}
+    # Will only add wires that are not present in used_wires yet, and to the end of used_wires
+    used_and_work_wires = used_wires | {
+        wire: None for op in tape for wire in getattr(op, "work_wires", [])
+    }
+    full_wire_map = {wire: ind for ind, wire in enumerate(used_and_work_wires)}
+    return full_wire_map, used_wire_map
 
 
 def default_bit_map(tape):
@@ -83,20 +90,29 @@ def convert_wire_order(tape, wire_order=None, show_all_wires=False):
             or only include ones used by operations in ``ops``
 
     Returns:
-        dict: map from wire labels to sequential positive integers
+        tuple[dict]: Two maps from wire labels to sequential positive integers. The first map
+        includes work wires, the second map excludes work wires.
     """
-    default = default_wire_map(tape)
+    full_wire_map, used_wire_map = default_wire_map(tape)
 
     if wire_order is None:
-        return default
+        # If no external wire order is dictated, the tape ordering is all we need to consider
+        return full_wire_map, used_wire_map
 
-    wire_order = list(wire_order) + [wire for wire in default if wire not in wire_order]
+    # Create wire order complemented by all wires in the tape mapping that are not in the order yet
+    full_wire_order = list(wire_order) + [wire for wire in full_wire_map if wire not in wire_order]
+    used_wire_order = list(wire_order) + [wire for wire in used_wire_map if wire not in wire_order]
 
     if not show_all_wires:
-        used_wires = {wire for op in tape for wire in op.wires}
-        wire_order = [wire for wire in wire_order if wire in used_wires]
+        # Filter out wires that are in wire_order but not in full_wire_map/used_wire_map
+        full_wire_order = [wire for wire in full_wire_order if wire in full_wire_map]
+        used_wire_order = [wire for wire in used_wire_order if wire in used_wire_map]
+
+    # Create consecutive integer mapping from ordered list
+    full_wire_map = {wire: ind for ind, wire in enumerate(full_wire_order)}
+    used_wire_map = {wire: ind for ind, wire in enumerate(used_wire_order)}
 
-    return {wire: ind for ind, wire in enumerate(wire_order)}
+    return full_wire_map, used_wire_map
 
 
 def unwrap_controls(op):


Test patch:
diff --git a/tests/drawer/test_drawer_utils.py b/tests/drawer/test_drawer_utils.py
index 179e0c308..8daafd48b 100644
--- a/tests/drawer/test_drawer_utils.py
+++ b/tests/drawer/test_drawer_utils.py
@@ -34,24 +34,40 @@ class TestDefaultWireMap:
     def test_empty(self):
         """Test creating an empty wire map"""
 
-        wire_map = default_wire_map([])
-        assert wire_map == {}
+        full_wire_map, used_wire_map = default_wire_map([])
+        assert full_wire_map == used_wire_map == {}
 
     def test_simple(self):
         """Test creating a wire map with wires that do not have successive ordering"""
 
         ops = [qml.PauliX(0), qml.PauliX(2), qml.PauliX(1)]
 
-        wire_map = default_wire_map(ops)
-        assert wire_map == {0: 0, 2: 1, 1: 2}
+        full_wire_map, used_wire_map = default_wire_map(ops)
+        assert full_wire_map == used_wire_map == {0: 0, 2: 1, 1: 2}
 
     def test_string_wires(self):
         """Test wire map works with string labelled wires."""
 
         ops = [qml.PauliY("a"), qml.CNOT(wires=("b", "c"))]
 
-        wire_map = default_wire_map(ops)
-        assert wire_map == {"a": 0, "b": 1, "c": 2}
+        full_wire_map, used_wire_map = default_wire_map(ops)
+        assert full_wire_map == used_wire_map == {"a": 0, "b": 1, "c": 2}
+
+    def test_work_wires(self):
+        """Test wire map works with string work wires, leading to a difference
+        between full_wire_map and used_wire_map."""
+
+        ops = [qml.PauliY("a"), qml.MultiControlledX(["b", 0, 9, 4], work_wires=[1, 5, "a"])]
+
+        full_wire_map, used_wire_map = default_wire_map(ops)
+        assert full_wire_map == {"a": 0, "b": 1, 0: 2, 9: 3, 4: 4, 1: 5, 5: 6}
+        assert used_wire_map == {"a": 0, "b": 1, 0: 2, 9: 3, 4: 4}
+
+        ops = [qml.MultiControlledX(["b", 0, 9, 4], work_wires=[1, 5, "a"]), qml.PauliY("a")]
+        full_wire_map, used_wire_map = default_wire_map(ops)
+        # Work-only wires always come after used wires
+        assert full_wire_map == {"b": 0, 0: 1, 9: 2, 4: 3, "a": 4, 1: 5, 5: 6}
+        assert used_wire_map == {"b": 0, 0: 1, 9: 2, 4: 3, "a": 4}
 
 
 class TestDefaultBitMap:
@@ -88,9 +104,8 @@ class TestConvertWireOrder:
 
         ops = [qml.PauliX(0), qml.PauliX(2), qml.PauliX(1)]
 
-        wire_map = convert_wire_order(ops)
-
-        assert wire_map == {0: 0, 2: 1, 1: 2}
+        full_wire_map, used_wire_map = convert_wire_order(ops)
+        assert full_wire_map == used_wire_map == {0: 0, 2: 1, 1: 2}
 
     def test_wire_order_ints(self):
         """Tests wire map produced when initial wires are integers."""
@@ -98,8 +113,8 @@ class TestConvertWireOrder:
         ops = [qml.PauliX(0), qml.PauliX(2), qml.PauliX(1)]
         wire_order = [2, 1, 0]
 
-        wire_map = convert_wire_order(ops, wire_order)
-        assert wire_map == {2: 0, 1: 1, 0: 2}
+        full_wire_map, used_wire_map = convert_wire_order(ops, wire_order)
+        assert full_wire_map == used_wire_map == {2: 0, 1: 1, 0: 2}
 
     def test_wire_order_str(self):
         """Test wire map produced when initial wires are strings."""
@@ -107,8 +122,8 @@ class TestConvertWireOrder:
         ops = [qml.CNOT(wires=("a", "b")), qml.PauliX("c")]
         wire_order = ("c", "b", "a")
 
-        wire_map = convert_wire_order(ops, wire_order)
-        assert wire_map == {"c": 0, "b": 1, "a": 2}
+        full_wire_map, used_wire_map = convert_wire_order(ops, wire_order)
+        assert full_wire_map == used_wire_map == {"c": 0, "b": 1, "a": 2}
 
     def test_show_all_wires_false(self):
         """Test when `show_all_wires` is set to `False` only used wires are in the map."""
@@ -116,17 +131,43 @@ class TestConvertWireOrder:
         ops = [qml.PauliX("a"), qml.PauliY("c")]
         wire_order = ["a", "b", "c", "d"]
 
-        wire_map = convert_wire_order(ops, wire_order, show_all_wires=False)
-        assert wire_map == {"a": 0, "c": 1}
+        full_wire_map, used_wire_map = convert_wire_order(ops, wire_order, show_all_wires=False)
+        assert full_wire_map == used_wire_map == {"a": 0, "c": 1}
 
     def test_show_all_wires_true(self):
         """Test when `show_all_wires` is set to `True` everything in ``wire_order`` is included."""
 
-        ops = [qml.PauliX("a"), qml.PauliY("c")]
+        ops = [qml.X("a"), qml.PauliY("c")]
         wire_order = ["a", "b", "c", "d"]
 
-        wire_map = convert_wire_order(ops, wire_order, show_all_wires=True)
-        assert wire_map == {"a": 0, "b": 1, "c": 2, "d": 3}
+        full_wire_map, used_wire_map = convert_wire_order(ops, wire_order, show_all_wires=True)
+        assert full_wire_map == used_wire_map == {"a": 0, "b": 1, "c": 2, "d": 3}
+
+    def test_with_work_wires(self):
+        """Tests wire map produced when work wires are present."""
+
+        ops = [qml.X(0), qml.ctrl(qml.X(2), control=[3, 1, 5], work_wires=[4, 0]), qml.X(1)]
+
+        # The work-only wire 4 does not show up in used_wire_map
+        full_wire_map, used_wire_map = convert_wire_order(ops, None)
+        # Control wires are added before target wires
+        assert full_wire_map == {0: 0, 3: 1, 1: 2, 5: 3, 2: 4, 4: 5}
+        assert used_wire_map == {0: 0, 3: 1, 1: 2, 5: 3, 2: 4}
+
+    def test_with_work_wires_wire_order(self):
+        """Tests wire map produced when work wires are present."""
+
+        ops = [qml.X(0), qml.ctrl(qml.X(2), control=[3, 1, 5], work_wires=[4, 0]), qml.X(1)]
+
+        wire_order = [2, 1, 0, 4]
+        # If we set show_all_wires to False, the work-only wire 4 does not show in used_wire_map
+        full_wire_map, used_wire_map = convert_wire_order(ops, wire_order, show_all_wires=False)
+        assert full_wire_map == {2: 0, 1: 1, 0: 2, 4: 3, 3: 4, 5: 5}
+        assert used_wire_map == {2: 0, 1: 1, 0: 2, 3: 3, 5: 4}
+
+        # If we set show_all_wires to True, the work-only wire 4 also appears in used_wire_map
+        full_wire_map, used_wire_map = convert_wire_order(ops, wire_order, show_all_wires=True)
+        assert full_wire_map == used_wire_map == {2: 0, 1: 1, 0: 2, 4: 3, 3: 4, 5: 5}
 
 
 class TestUnwrapControls: