Repository: PennyLaneAI/pennylane

Issue:
[BUG] Auxiliary wire handling in `metric_tensor` and `hadamard_grad`
### Expected behavior

The handling of the `aux_wire` and `device_wires` kwargs is the same for `transforms.metric_tensor` and `gradients.hadamard_grad`, and as many use cases as possible are supported.

### Actual behavior

Differences between the two functions:
1. `metric_tensor` can handle an `aux_wire` input without a `device_wires` input, `hadamard_grad` raises an error (which is not necessary imo).
2. `hadamard_grad` can handle `aux_wire` to be a `Wires` object, `metric_tensor` can not.
3. `metric_tensor` can handle `aux_wire` to be a multi-character string or a simple integer, `hadamard_grad` can not.
4. `hadamard_grad` correctly detects reusage of a wire from the tape, `metric_tensor` does not (and produces flawed circuits)

### Additional information

This can easily be fixed by unifying the handling of `aux_wire` and `device_wires` from the two functions and cherrypicking the best of both worlds.

### Source code

Take the circuit
```python
dev = qml.device("default.qubit", wires=[0, 1, 2, "aux"])

@qml.qnode(dev, interface="jax")
def circuit(x, argnums=None):
    qml.RX(x[0], 0)
    qml.CNOT([0, 1])
    qml.RX(x[1], 0)
    return qml.expval(Z(0))

x = jnp.array([1.3, 0.2])
```
Then
```python
# 1
qml.gradients.hadamard_grad(circuit, aux_wire=qml.wires.Wires(2))(x) # does not work, None not subscriptable
# 2
qml.transforms.metric_tensor(circuit, aux_wire=qml.wires.Wires(2))(x) # does not work, wire not recognized
# 3
qml.gradients.hadamard_grad(circuit, aux_wire=2)(x) # does not work, int not subscriptable
qml.gradients.hadamard_grad(circuit, aux_wire="aux")(x) # does not work, wire not found
# 4
qml.transforms.metric_tensor(circuit, aux_wire=1)(x) # runs through but creates invalid tapes, reusing tape wire.
qml.transforms.metric_tensor(circuit, aux_wire=0)(x) #raises not-so-useful error
```


### Tracebacks

_No response_

### System information

```shell
pl dev
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4328?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#4328](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4328?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (e09b0a3) into [master](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/34d2fb218f38d88d01f4b742a915f3149cd0c6ac?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (34d2fb2) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #4328   +/-   ##
=======================================
  Coverage   99.79%   99.79%           
=======================================
  Files         351      351           
  Lines       32101    32102    +1     
=======================================
+ Hits        32034    32035    +1     
  Misses         67       67           
```


| [Impacted Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4328?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/gradients/hadamard\_gradient.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4328?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9oYWRhbWFyZF9ncmFkaWVudC5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/metric\_tensor.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4328?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvbWV0cmljX3RlbnNvci5weQ==) | `99.61% <100.00%> (+0.01%)` | :arrow_up: |



Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 5ea039d3a..f48133724 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -6,6 +6,10 @@
 
 <h3>Improvements 🛠</h3>
 
+* Treat auxiliary wires and device wires in the same way in `transforms.metric_tensor`
+  as in `gradients.hadamard_grad`. Support all valid wire input formats for `aux_wire`.
+  [(#4328)](https://github.com/PennyLaneAI/pennylane/pull/4328)
+
 * `qml.equal` no longer raises errors when operators or measurements of different types are compared.
   Instead, it returns `False`.
   [(#4315)](https://github.com/PennyLaneAI/pennylane/pull/4315)
@@ -76,6 +80,10 @@
 <h3>Documentation 📝</h3>
 
 <h3>Bug fixes 🐛</h3>
+  
+* Stop `metric_tensor` from accidentally catching errors that stem from
+  flawed wires assignments in the original circuit, leading to recursion errors
+  [(#4328)](https://github.com/PennyLaneAI/pennylane/pull/4328)
 
 * Raise a warning if control indicators are hidden when calling `qml.draw_mpl`
   [(#4295)](https://github.com/PennyLaneAI/pennylane/pull/4295)
@@ -91,4 +99,5 @@ Edward Jiang,
 Christina Lee,
 Mudit Pandey,
 Borja Requena,
-Matthew Silverman
+Matthew Silverman,
+David Wierichs,
diff --git a/pennylane/gradients/hadamard_gradient.py b/pennylane/gradients/hadamard_gradient.py
index bad143fbf..9dabfbf21 100644
--- a/pennylane/gradients/hadamard_gradient.py
+++ b/pennylane/gradients/hadamard_gradient.py
@@ -196,18 +196,8 @@ def _hadamard_grad(
 
     argnum = [i for i, dm in method_map.items() if dm == "A"]
 
-    if device_wires and len(tape.wires) == len(device_wires):
-        raise qml.QuantumFunctionError("The device has no free wire for the auxiliary wire.")
-
-    # Get default for aux_wire
-    if aux_wire is None:
-        aux_wire = _get_aux_wire(aux_wire, tape, device_wires)
-    elif aux_wire[0] in tape.wires:
-        raise qml.QuantumFunctionError("The auxiliary wire is already used.")
-    elif aux_wire[0] not in device_wires:
-        raise qml.QuantumFunctionError(
-            "The requested auxiliary wire does not exist on the used device."
-        )
+    # Validate or get default for aux_wire
+    aux_wire = _get_aux_wire(aux_wire, tape, device_wires)
 
     g_tapes, processing_fn = _expval_hadamard_grad(tape, argnum, aux_wire)
 
diff --git a/pennylane/transforms/metric_tensor.py b/pennylane/transforms/metric_tensor.py
index 77fe60a49..fb33f00ea 100644
--- a/pennylane/transforms/metric_tensor.py
+++ b/pennylane/transforms/metric_tensor.py
@@ -86,9 +86,10 @@ def metric_tensor(
         allow_nonunitary (bool): Whether non-unitary operations are allowed in circuits
             created by the transform. Only relevant if ``approx`` is ``None``.
             Should be set to ``True`` if possible to reduce cost.
-        aux_wire (int or str or pennylane.wires.Wires): Auxiliary wire to be used for
-            Hadamard tests. If ``None`` (the default), a suitable wire is inferred
-            from the (number of) used wires in the original circuit and ``device_wires``.
+        aux_wire (None or int or str or Sequence or pennylane.wires.Wires): Auxiliary wire to
+            be used for Hadamard tests. If ``None`` (the default), a suitable wire is inferred
+            from the (number of) used wires in the original circuit and ``device_wires``,
+            if the latter are given.
         device_wires (.wires.Wires): Wires of the device that is going to be used for the
             metric tensor. Facilitates finding a default for ``aux_wire`` if ``aux_wire``
             is ``None``.
@@ -424,22 +425,29 @@ def qnode_execution_wrapper(self, qnode, targs, tkwargs):
         try:
             mt = mt_fn(*args, **kwargs)
         except qml.wires.WireError as e:
-            if str(e) == "No device wires are unused by the tape.":
+            revert_text = (
+                "\n\nReverting to the block-diagonal approximation. It will often be "
+                "much more efficient to request the block-diagonal approximation directly!"
+            )
+            other_mt_errors = [
+                "The requested auxiliary wire is already in use by the circuit.",
+                "The requested auxiliary wire does not exist on the used device.",
+            ]
+
+            if str(e) == "The device has no free wire for the auxiliary wire.":
                 warnings.warn(
-                    "The device does not have a wire that is not used by the tape."
-                    "\n\nReverting to the block-diagonal approximation. It will often be "
-                    "much more efficient to request the block-diagonal approximation directly!"
+                    "The device does not have a wire that is not used by the circuit." + revert_text
                 )
-            else:
+            elif str(e) in other_mt_errors:
                 warnings.warn(
                     "An auxiliary wire is not available."
                     "\n\nThis can occur when computing the full metric tensor via the "
                     "Hadamard test, and the device does not provide an "
                     "additional wire or the requested auxiliary wire does not exist "
-                    "on the device."
-                    "\n\nReverting to the block-diagonal approximation. It will often be "
-                    "much more efficient to request the block-diagonal approximation directly!"
+                    "on the device." + revert_text
                 )
+            else:
+                raise e
             tkwargs["approx"] = "block-diag"
             return self(qnode, *targs, **tkwargs)(*args, **kwargs)
 
@@ -853,25 +861,36 @@ def _get_aux_wire(aux_wire, tape, device_wires):
     r"""Determine an unused wire to be used as auxiliary wire for Hadamard tests.
 
     Args:
-        aux_wire (object): Input auxiliary wire. Returned unmodified if not ``None``
+        aux_wire (object): Input auxiliary wire. May be one of a variety of input formats:
+            If ``None``, try to infer a reasonable choice based on the number of wires used
+            in the ``tape``, and based on ``device_wires``, if they are not ``None``.
+            If an ``int``, a ``str`` or a ``Sequence``, convert the input to a ``Wires``
+            object and take the first entry of the result. This leads to consistent behaviour
+            between ``_get_aux_wire`` and the ``Wires`` class.
+            If a ``Wires`` instance already, the conversion to such an instance is performed
+            trivially as well (also see the source code of ``~.Wires``).
         tape (pennylane.tape.QuantumTape): Tape to infer the wire for
         device_wires (.wires.Wires): Wires of the device that is going to be used for the
             metric tensor. Facilitates finding a default for ``aux_wire`` if ``aux_wire``
             is ``None`` .
 
     Returns:
-        object: The auxiliary wire to be used. Equals ``aux_wire`` if it was not ``None`` ,
+        object: The auxiliary wire to be used. Equals ``aux_wire`` if it was not ``None``\ ,
         and an often reasonable choice else.
     """
     if aux_wire is not None:
+        aux_wire = qml.wires.Wires(aux_wire)[0]
+        if aux_wire in tape.wires:
+            msg = "The requested auxiliary wire is already in use by the circuit."
+            raise qml.wires.WireError(msg)
         if device_wires is None or aux_wire in device_wires:
             return aux_wire
-        raise qml.wires.WireError("The requested aux_wire does not exist on the used device.")
+        raise qml.wires.WireError("The requested auxiliary wire does not exist on the used device.")
 
     if device_wires is not None:
+        if len(device_wires) == len(tape.wires):
+            raise qml.wires.WireError("The device has no free wire for the auxiliary wire.")
         unused_wires = qml.wires.Wires(device_wires.toset().difference(tape.wires.toset()))
-        if not unused_wires:
-            raise qml.wires.WireError("No device wires are unused by the tape.")
         return unused_wires[0]
 
     _wires = tape.wires


Test patch:
diff --git a/tests/gradients/core/test_hadamard_gradient.py b/tests/gradients/core/test_hadamard_gradient.py
index 750ee1734..987c495dc 100644
--- a/tests/gradients/core/test_hadamard_gradient.py
+++ b/tests/gradients/core/test_hadamard_gradient.py
@@ -544,7 +544,7 @@ class TestHadamardGradEdgeCases:
     device_wires = [qml.wires.Wires([0, 1, "aux"])]
     device_wires_no_aux = [qml.wires.Wires([0, 1, 2])]
 
-    working_wires = [None, qml.wires.Wires("aux")]
+    working_wires = [None, qml.wires.Wires("aux"), "aux"]
     already_used_wires = [qml.wires.Wires(0), qml.wires.Wires(1)]
 
     @pytest.mark.parametrize("aux_wire", working_wires)
@@ -566,6 +566,8 @@ class TestHadamardGradEdgeCases:
 
         tapes, _ = qml.gradients.hadamard_grad(tape, aux_wire=aux_wire, device_wires=dev.wires)
         assert len(tapes) == 2
+        tapes, _ = qml.gradients.hadamard_grad(tape, aux_wire=aux_wire)
+        assert len(tapes) == 2
 
     @pytest.mark.parametrize("aux_wire", already_used_wires)
     @pytest.mark.parametrize("device_wires", device_wires)
@@ -583,7 +585,8 @@ class TestHadamardGradEdgeCases:
 
         tape = qml.tape.QuantumScript.from_queue(q)
 
-        with pytest.raises(qml.QuantumFunctionError, match="The auxiliary wire is already."):
+        _match = "The requested auxiliary wire is already in use by the circuit"
+        with pytest.raises(qml.wires.WireError, match=_match):
             qml.gradients.hadamard_grad(tape, aux_wire=aux_wire, device_wires=dev.wires)
 
     @pytest.mark.parametrize("device_wires", device_wires_no_aux)
@@ -601,30 +604,36 @@ class TestHadamardGradEdgeCases:
             qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        with pytest.raises(
-            qml.QuantumFunctionError,
-            match="The requested auxiliary wire does not exist on the used device.",
-        ):
+        _match = "The requested auxiliary wire does not exist on the used device"
+        with pytest.raises(qml.wires.WireError, match=_match):
             qml.gradients.hadamard_grad(tape, aux_wire=aux_wire, device_wires=dev.wires)
 
-    @pytest.mark.parametrize("aux_wire", working_wires + already_used_wires)
+    @pytest.mark.parametrize("aux_wire", [None] + already_used_wires)
     def test_device_not_enough_wires(self, aux_wire):
-        """Test that an error is raised when the device cannot accept an auxiliary wire because it is full."""
+        """Test that an error is raised when the device cannot accept an auxiliary wire
+        because it is full."""
         dev = qml.device("default.qubit", wires=2)
-        x = 0.543
-        y = -0.654
 
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.RX(x, wires=[0])
-            qml.RY(y, wires=[1])
-            qml.CNOT(wires=[0, 1])
-            qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+        m = qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+        tape = qml.tape.QuantumScript([qml.RX(0.543, wires=[0]), qml.RY(-0.654, wires=[1])], [m])
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        if aux_wire is None:
+            _match = "The device has no free wire for the auxiliary wire."
+        else:
+            _match = "The requested auxiliary wire is already in use by the circuit."
+        with pytest.raises(qml.wires.WireError, match=_match):
+            qml.gradients.hadamard_grad(tape, aux_wire=aux_wire, device_wires=dev.wires)
 
-        with pytest.raises(
-            qml.QuantumFunctionError, match="The device has no free wire for the auxiliary wire."
-        ):
+    def test_device_wire_does_not_exist(self):
+        """Test that an error is raised when the device cannot accept an auxiliary wire
+        because it does not exist on the device."""
+        aux_wire = qml.wires.Wires("aux")
+        dev = qml.device("default.qubit", wires=2)
+        m = qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+        tape = qml.tape.QuantumScript([qml.RX(0.543, wires=[0]), qml.RY(-0.654, wires=[1])], [m])
+
+        _match = "The requested auxiliary wire does not exist on the used device."
+        with pytest.raises(qml.wires.WireError, match=_match):
             qml.gradients.hadamard_grad(tape, aux_wire=aux_wire, device_wires=dev.wires)
 
     def test_empty_circuit(self):
diff --git a/tests/legacy/test_legacy_metric_tensor_old.py b/tests/legacy/test_legacy_metric_tensor_old.py
index 73294757e..0a051fb94 100644
--- a/tests/legacy/test_legacy_metric_tensor_old.py
+++ b/tests/legacy/test_legacy_metric_tensor_old.py
@@ -833,7 +833,7 @@ class TestMetricTensor:
 
         error_msg = (
             "Some parameters specified in argnum are not in the "
-            "trainable parameters \[0, 1, 2, 3\] of the tape "
+            r"trainable parameters \[0, 1, 2, 3\] of the tape "
             "and will be ignored. This may be caused by attempting to "
             "differentiate with respect to parameters that are not marked "
             "as trainable."
@@ -1758,9 +1758,13 @@ def test_get_aux_wire_with_device_wires():
     tape = qml.tape.QuantumScript.from_queue(q)
     device_wires = qml.wires.Wires([0, "aux", "one"])
 
-    assert _get_aux_wire(0, tape, device_wires) == 0
-    assert _get_aux_wire("one", tape, device_wires) == "one"
     assert _get_aux_wire(None, tape, device_wires) == "aux"
+    assert _get_aux_wire("aux", tape, device_wires) == "aux"
+    _match = "The requested auxiliary wire is already in use by the circuit."
+    with pytest.raises(qml.wires.WireError, match=_match):
+        _get_aux_wire("one", tape, device_wires)
+    with pytest.raises(qml.wires.WireError, match=_match):
+        _get_aux_wire(0, tape, device_wires)
 
 
 def test_get_aux_wire_with_unavailable_aux():
@@ -1771,5 +1775,5 @@ def test_get_aux_wire_with_unavailable_aux():
         qml.RX(x, wires="one")
     tape = qml.tape.QuantumScript.from_queue(q)
     device_wires = qml.wires.Wires([0, "one"])
-    with pytest.raises(qml.wires.WireError, match="The requested aux_wire does not exist"):
+    with pytest.raises(qml.wires.WireError, match="The requested auxiliary wire does not exist"):
         _get_aux_wire("two", tape, device_wires)
diff --git a/tests/transforms/test_metric_tensor.py b/tests/transforms/test_metric_tensor.py
index c4b4e78d9..c28fcb236 100644
--- a/tests/transforms/test_metric_tensor.py
+++ b/tests/transforms/test_metric_tensor.py
@@ -1745,6 +1745,25 @@ def test_no_error_missing_aux_wire_not_used(recwarn):
     assert len(recwarn) == 0
 
 
+def test_raises_circuit_that_uses_missing_wire():
+    """Test that an error in the original circuit is reraised properly and not caught. This avoids
+    accidentally catching relevant errors, which can lead to a recursion error."""
+
+    dev = qml.device("default.qubit", wires=[0, "b"])
+
+    @qml.qnode(dev)
+    def circuit(x):
+        """Flawed circuit that uses a wire which is not on the device."""
+        qml.RX(x[0], 0)
+        qml.CNOT([0, 1])  # wire 1 is not on the device
+        qml.RX(x[1], 0)
+        return qml.expval(qml.PauliZ(0))
+
+    x = np.array([1.3, 0.2])
+    with pytest.raises(qml.wires.WireError, match=r"Did not find some of the wires \(0, 1\)"):
+        qml.transforms.metric_tensor(circuit)(x)
+
+
 def aux_wire_ansatz_0(x, y):
     qml.RX(x, wires=0)
     qml.RY(y, wires=2)
@@ -1781,9 +1800,13 @@ def test_get_aux_wire_with_device_wires():
     tape = qml.tape.QuantumScript.from_queue(q)
     device_wires = qml.wires.Wires([0, "aux", "one"])
 
-    assert _get_aux_wire(0, tape, device_wires) == 0
-    assert _get_aux_wire("one", tape, device_wires) == "one"
     assert _get_aux_wire(None, tape, device_wires) == "aux"
+    assert _get_aux_wire("aux", tape, device_wires) == "aux"
+    _match = "The requested auxiliary wire is already in use by the circuit."
+    with pytest.raises(qml.wires.WireError, match=_match):
+        _get_aux_wire("one", tape, device_wires)
+    with pytest.raises(qml.wires.WireError, match=_match):
+        _get_aux_wire(0, tape, device_wires)
 
 
 def test_get_aux_wire_with_unavailable_aux():
@@ -1794,5 +1817,5 @@ def test_get_aux_wire_with_unavailable_aux():
         qml.RX(y, wires="one")
     tape = qml.tape.QuantumScript.from_queue(q)
     device_wires = qml.wires.Wires([0, "one"])
-    with pytest.raises(qml.wires.WireError, match="The requested aux_wire does not exist"):
+    with pytest.raises(qml.wires.WireError, match="The requested auxiliary wire does not exist"):
         _get_aux_wire("two", tape, device_wires)