Repository: PennyLaneAI/pennylane

Issue:
[BUG] `qml.map_wires` does not work correctly with `qml.cond`
### Expected behavior

The mid-circuit measurements that a `qml.cond` depends on are mapped to the new wires correctly.

### Actual behavior

The wires of the measurements used to control the `qml.cond` are not mapped.

### Additional information

_No response_

### Source code

```shell
with qml.queuing.AnnotatedQueue() as q:
    m = qml.measure("a")
    qml.cond(m, qml.PauliZ)("b")

qs = qml.tape.QuantumScript.from_queue(q)
[tape], _ = qml.map_wires(qs, {"a": 0, "b": 1})
tape.operations[1].meas_val.measurements
```


### Tracebacks

```shell
[measure(wires=['a'])]
```


### System information

```shell
Name: PennyLane
Version: 0.34.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/mudit.pandey/repos/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing_extensions
Required-by: PennyLane-Lightning

Platform info:           macOS-13.6.1-arm64-arm-64bit
Python version:          3.10.12
Numpy version:           1.26.1
Scipy version:           1.11.1
Installed devices:
- default.gaussian (PennyLane-0.34.0.dev0)
- default.mixed (PennyLane-0.34.0.dev0)
- default.qubit (PennyLane-0.34.0.dev0)
- default.qubit.autograd (PennyLane-0.34.0.dev0)
- default.qubit.jax (PennyLane-0.34.0.dev0)
- default.qubit.legacy (PennyLane-0.34.0.dev0)
- default.qubit.tf (PennyLane-0.34.0.dev0)
- default.qubit.torch (PennyLane-0.34.0.dev0)
- default.qutrit (PennyLane-0.34.0.dev0)
- null.qubit (PennyLane-0.34.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.33.1)
```


### Existing GitHub issues

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


Discussion:

Resolving patch:
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -127,6 +127,8 @@ class MeasurementProcess(ABC):
             where the instance has to be identified
     """
 
+    # pylint:disable=too-many-instance-attributes
+
     def __init_subclass__(cls, **_):
         register_pytree(cls, cls._flatten, cls._unflatten)
 
@@ -469,10 +471,8 @@ def map_wires(self, wire_map: dict):
         """
         new_measurement = copy.copy(self)
         if self.mv is not None:
-            mv = copy.copy(self.mv)
-            mv.measurements = [m.map_wires(wire_map=wire_map) for m in mv.measurements]
-            new_measurement.mv = mv
-        if self.obs is not None:
+            new_measurement.mv = self.mv.map_wires(wire_map=wire_map)
+        elif self.obs is not None:
             new_measurement.obs = self.obs.map_wires(wire_map=wire_map)
         elif self._wires is not None:
             new_measurement._wires = Wires([wire_map.get(wire, wire) for wire in self.wires])
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -305,6 +305,19 @@ def branches(self):
             ret_dict[branch] = self.processing_fn(*branch)
         return ret_dict
 
+    def map_wires(self, wire_map):
+        """Returns a copy of the current ``MeasurementValue`` with the wires of each measurement changed
+        according to the given wire map.
+
+        Args:
+            wire_map (dict): dictionary containing the old wires as keys and the new wires as values
+
+        Returns:
+            MeasurementValue: new ``MeasurementValue`` instance with measurement wires mapped
+        """
+        mapped_measurements = [m.map_wires(wire_map) for m in self.measurements]
+        return MeasurementValue(mapped_measurements, self.processing_fn)
+
     def _transform_bin_op(self, base_bin, other):
         """Helper function for defining dunder binary operations."""
         if isinstance(other, MeasurementValue):
diff --git a/pennylane/ops/op_math/condition.py b/pennylane/ops/op_math/condition.py
--- a/pennylane/ops/op_math/condition.py
+++ b/pennylane/ops/op_math/condition.py
@@ -55,6 +55,11 @@ def __init__(self, expr, then_op: Type[Operation], id=None):
     def label(self, decimals=None, base_label=None, cache=None):
         return self.then_op.label(decimals=decimals, base_label=base_label, cache=cache)
 
+    def map_wires(self, wire_map):
+        meas_val = self.meas_val.map_wires(wire_map)
+        then_op = self.then_op.map_wires(wire_map)
+        return Conditional(meas_val, then_op=then_op)
+
 
 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_drawable_layers.py b/tests/drawer/test_drawable_layers.py
--- a/tests/drawer/test_drawable_layers.py
+++ b/tests/drawer/test_drawable_layers.py
@@ -138,6 +138,26 @@ def test_all_wires_measurement(self, measurement):
 
         assert layers == [[ops[0]], [ops[1]], [ops[2]]]
 
+    def test_mid_measure_custom_wires(self):
+        """Test that custom wires do not break the drawing of mid-circuit measurements."""
+
+        def teleport(state):
+            qml.QubitStateVector(state, wires=["A"])
+            qml.Hadamard(wires="a")
+            qml.CNOT(wires=["a", "B"])
+            qml.CNOT(wires=["A", "a"])
+            qml.Hadamard(wires="A")
+            m0 = qml.measure("A")
+            m1 = qml.measure("a")
+            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)
+        assert layers == [ops[:2]] + [[op] for op in ops[2:]]
+
 
 class TestMidMeasure:
     """Tests the various changes from mid-circuit measurements."""
diff --git a/tests/measurements/test_mid_measure.py b/tests/measurements/test_mid_measure.py
--- a/tests/measurements/test_mid_measure.py
+++ b/tests/measurements/test_mid_measure.py
@@ -388,6 +388,14 @@ def test_complex_str(self):
 if m0=1,m1=1 => 2"""
         )
 
+    def test_map_wires(self):
+        """Test that map_wires works as expected."""
+        a = MeasurementValue([mp1], lambda v: v)
+        b = a.map_wires({0: "b"})
+        [new_meas] = b.measurements
+        assert new_meas.wires == Wires(["b"])
+        assert new_meas.id == mp1.id
+
 
 unary_dunders = ["__invert__"]
 
diff --git a/tests/ops/op_math/test_condition.py b/tests/ops/op_math/test_condition.py
--- a/tests/ops/op_math/test_condition.py
+++ b/tests/ops/op_math/test_condition.py
@@ -190,10 +190,13 @@ def g():
             m_0 = qml.measure(1)
             qml.cond(m_0, g, f)()  # Check that the same error is raised when f and g are swapped
 
+
+class TestAdditionalCond:
+    """Test additional/misc functionality relating to qml.cond"""
+
     @pytest.mark.parametrize("inp", [1, "string", qml.PauliZ(0)])
-    def test_cond_error_unrecognized_input(self, inp, terminal_measurement):
+    def test_cond_error_unrecognized_input(self, inp):
         """Test that an error is raised when the input is not recognized."""
-        # pylint: disable=unused-argument
 
         with pytest.raises(
             ConditionalTransformError,
@@ -202,6 +205,17 @@ def test_cond_error_unrecognized_input(self, inp, terminal_measurement):
             m_0 = qml.measure(1)
             qml.cond(m_0, inp)()
 
+    def test_map_wires(self):
+        """Tests the cond.map_wires function."""
+        with qml.queuing.AnnotatedQueue() as q:
+            m_0 = qml.measure(0)
+            qml.cond(m_0, qml.PauliX)(1)
+
+        meas, cond_op = q.queue
+        mapped_cond = cond_op.map_wires({0: "a", 1: "b"})
+        assert mapped_cond.meas_val.measurements == [meas.map_wires({0: "a"})]
+        assert mapped_cond.then_op == qml.PauliX("b")
+
 
 @pytest.mark.parametrize("op_class", [qml.PauliY, qml.Toffoli, qml.Hadamard, qml.CZ])
 def test_conditional_label(op_class):