Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] metric_tensor doesn't work with TrotterProduct with more than one Trotter step
### Expected behavior

The question originates in [Forum thread 8118](https://discuss.pennylane.ai/t/issue-with-qfim-for-time-evolved-mixed-state/8118), where the goal is to calculate the quantum Fisher information matrix for a circuit using `qml.ApproxTimeEvolution`. The expectation is to be able to calculate this with no errors.

### Actual behavior

However, both when using `qml.ApproxTimeEvolution` or `qml.TrotterProduct` the metric_tensor fails and raises an error.

### Additional information

I'm not sure if this is expected or not, but in case it's expected the error should be more clear and/or the docs for metric_tensor should be updated to include the limitations for using it. I used a QuantumScript and TrotterProduct to try to narrow down the issue in the example below, but the problem can also be observed with the MWE posted in the [corresponding Forum thread (on post 19)](https://discuss.pennylane.ai/t/issue-with-qfim-for-time-evolved-mixed-state/8118/19). 

Edit: there seems to be a workaround. It's in the Forum thread (post 20).

### Source code

```shell
import pennylane as qml
from pennylane import numpy as pnp
from pennylane.tape import QuantumScript

# Parameter values
theta = pnp.array([pnp.pi / 4], requires_grad=True)

coeffs = [1, 1]
op = [qml.X(0), qml.X(1)]
H_Z = qml.Hamiltonian(coeffs, op)

ops = [qml.TrotterProduct(H_Z, theta, order=1, n=2)] # mt_fn works with n=1 but not with n=2

qscript = QuantumScript(ops, [qml.expval(qml.Z(0))])


mt_fn = qml.metric_tensor(qscript) # mt_fn works with n=1 but not with n=2 in TrotterProduct
```

### Tracebacks

```shell
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipython-input-4-2498340384.py in <cell line: 0>()
     18 
     19 
---> 20 mt_fn = qml.metric_tensor(qscript) # mt_fn works with n=1 but not with n=2 in TrotterProduct
     21 
     22 #metric_tensor_cov_matrix(qscript, argnum=[0], diag_approx=False)

6 frames
/usr/local/lib/python3.11/dist-packages/pennylane/transforms/core/transform_dispatcher.py in __call__(self, *targs, **tkwargs)
    239                 start = 0
    240                 for tape in expanded_tapes:
--> 241                     intermediate_tapes, post_processing_fn = self._transform(
    242                         tape, *targs, **tkwargs
    243                     )

/usr/local/lib/python3.11/dist-packages/pennylane/gradients/metric_tensor.py in metric_tensor(tape, argnum, approx, allow_nonunitary, aux_wire, device_wires)
    407 
    408     if approx is None:
--> 409         tapes, processing_fn = _metric_tensor_hadamard(
    410             tape, argnum, allow_nonunitary, aux_wire, device_wires
    411         )

/usr/local/lib/python3.11/dist-packages/pennylane/gradients/metric_tensor.py in _metric_tensor_hadamard(tape, argnum, allow_nonunitary, aux_wire, device_wires)
    691         layer_ids,
    692         obs_ids,
--> 693     ) = _metric_tensor_cov_matrix(tape, argnum, diag_approx=False)
    694 
    695     # Obtain layers of parametrized operations and account for the discrepancy between trainable

/usr/local/lib/python3.11/dist-packages/pennylane/gradients/metric_tensor.py in _metric_tensor_cov_matrix(tape, argnum, diag_approx)
    473 
    474     i = 0
--> 475     for queue, curr_ops, param_idx, _ in graph.iterate_parametrized_layers():
    476         params_list.append(param_idx)
    477         in_argnum_list.append([p in argnum for p in param_idx])

/usr/local/lib/python3.11/dist-packages/pennylane/circuit_graph.py in iterate_parametrized_layers(self)
    387         """
    388         # iterate through each layer
--> 389         for ops, param_inds in self.parametrized_layers:
    390             pre_queue = self.ancestors_in_order(ops)
    391             post_queue = self.descendants_in_order(ops)

/usr/local/lib/python3.11/dist-packages/pennylane/circuit_graph.py in parametrized_layers(self)
    365 
    366                 # get all predecessor ops of the op
--> 367                 sub = self.ancestors((op,))
    368 
    369                 # check if any of the dependents are in the

/usr/local/lib/python3.11/dist-packages/pennylane/circuit_graph.py in ancestors(self, ops, sort)
    262             )
    263         if any(len(self._inds_for_objs[WrappedObj(op)]) > 1 for op in ops):
--> 264             raise ValueError(
    265                 "Cannot calculate ancestors for an operator that occurs multiple times."
    266             )

ValueError: Cannot calculate ancestors for an operator that occurs multiple times.
```

### System information

```shell
Name: pennylane
Version: 0.42.0.dev56
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: 
Author: 
Author-email: 
License: 
Location: /usr/local/lib/python3.11/dist-packages
Requires: appdirs, autograd, autoray, cachetools, diastatic-malt, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, tomlkit, typing_extensions
Required-by: PennyLane_Lightning

Platform info:           Linux-6.1.123+-x86_64-with-glibc2.35
Python version:          3.11.13
Numpy version:           2.0.2
Scipy version:           1.15.3
Installed devices:
- default.clifford (pennylane-0.42.0.dev56)
- default.gaussian (pennylane-0.42.0.dev56)
- default.mixed (pennylane-0.42.0.dev56)
- default.qubit (pennylane-0.42.0.dev56)
- default.qutrit (pennylane-0.42.0.dev56)
- default.qutrit.mixed (pennylane-0.42.0.dev56)
- default.tensor (pennylane-0.42.0.dev56)
- null.qubit (pennylane-0.42.0.dev56)
- reference.qubit (pennylane-0.42.0.dev56)
- lightning.qubit (PennyLane_Lightning-0.41.1)
```

### Existing GitHub issues

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

Pull request description:
Metric tensor trotter fix
**Context:**

When `TrotterProduct` was in the circuit, `metric_tensor` errored out:

```
import pennylane as qml
from pennylane import numpy as pnp
from pennylane.tape import QuantumScript

# Parameter values
theta = pnp.array([pnp.pi / 4], requires_grad=True)

coeffs = [1, 1]
op = [qml.X(0), qml.X(1)]
H_Z = qml.Hamiltonian(coeffs, op)

ops = [qml.TrotterProduct(H_Z, theta, order=1, n=2)] # mt_fn works with n=1 but not with n=2

qscript = QuantumScript(ops, [qml.expval(qml.Z(0))])

mt_fn = qml.metric_tensor(qscript) # mt_fn works with n=1 but not with n=2 in TrotterProduct
```
```
ValueError: Cannot calculate ancestors for an operator that occurs multiple times.
```

**Description of the Change:**

Updates `parametrized_layers` to calculate ancestors and descendants based on the operator index, rather than trying to find the operator's position in the circuit.

**Benefits:**

Allows for the same operator instance to exist multiple times.

**Possible Drawbacks:**

**Related GitHub Issues:**

Fixes #7717  [sc-93784]


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index c021b70e3..58fb73553 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -1076,6 +1076,10 @@ may move operations across a `Snapshot`.
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixes `CircuitGraph.iterate_parametrized_layers`, and thus `metric_tensor`, when the same operation occurs multiple
+  times in the circuit.
+  [(#7757)](https://github.com/PennyLaneAI/pennylane/pull/7757)
+
 * Fixes a bug with transforms that require the classical Jacobian applied to QNodes, where only
   some arguments are trainable and an intermediate transform does not preserve trainability information.
   [(#7345)](https://github.com/PennyLaneAI/pennylane/pull/7345)
diff --git a/pennylane/circuit_graph.py b/pennylane/circuit_graph.py
index f0b8dd8bc..61ba58690 100644
--- a/pennylane/circuit_graph.py
+++ b/pennylane/circuit_graph.py
@@ -17,7 +17,7 @@ representation of a quantum circuit from an Operator queue.
 """
 from collections import defaultdict, namedtuple
 from functools import cached_property
-from typing import List, Optional, Union
+from typing import Optional, Sequence, Union
 
 import numpy as np
 import rustworkx as rx
@@ -34,13 +34,14 @@ def _get_wires(obj, all_wires):
     return all_wires if len(obj.wires) == 0 else obj.wires
 
 
-Layer = namedtuple("Layer", ["ops", "param_inds"])
+Layer = namedtuple("Layer", ["ops", "param_inds", "ops_inds"])
 """Parametrized layer of the circuit.
 
 Args:
 
     ops (list[Operator]): parametrized operators in the layer
     param_inds (list[int]): corresponding free parameter indices
+    ops_inds (list[int]): the indices into the circuit for ops
 """
 # TODO define what a layer is
 
@@ -84,7 +85,7 @@ class CircuitGraph:
 
     Args:
         ops (Iterable[.Operator]): quantum operators constituting the circuit, in temporal order
-        obs (List[Union[MeasurementProcess, Operator]]): terminal measurements, in temporal order
+        obs (list[Union[MeasurementProcess, Operator]]): terminal measurements, in temporal order
         wires (.Wires): The addressable wire registers of the device that will be executing this graph
         par_info (Optional[list[dict]]): Parameter information. For each index, the entry is a dictionary containing an operation
         and an index into that operation's parameters.
@@ -97,7 +98,7 @@ class CircuitGraph:
     def __init__(
         self,
         ops: list[Union[Operator, MeasurementProcess]],
-        obs: List[Union[MeasurementProcess, Operator]],
+        obs: list[Union[MeasurementProcess, Operator]],
         wires: Wires,
         par_info: Optional[list[dict]] = None,
         trainable_params: Optional[set[int]] = None,
@@ -195,7 +196,7 @@ class CircuitGraph:
         Currently the topological order is determined by the queue index.
 
         Returns:
-            List[Union[MeasurementProcess, Operator]]: observables
+            list[Union[MeasurementProcess, Operator]]: observables
         """
         return self._observables
 
@@ -263,6 +264,7 @@ class CircuitGraph:
         if any(len(self._inds_for_objs[WrappedObj(op)]) > 1 for op in ops):
             raise ValueError(
                 "Cannot calculate ancestors for an operator that occurs multiple times."
+                "Please use ancestors_of_indexes instead."
             )
         ancestors = set()
         for op in ops:
@@ -273,6 +275,36 @@ class CircuitGraph:
             ancestors = sorted(ancestors)
         return [self._queue[ind] for ind in ancestors]
 
+    def ancestors_of_indexes(self, indexes: Sequence[int], sort=False):
+        """Ancestors of a given set of operators.
+
+        Args:
+            indexes (Sequence[int]) : the index into the queue for the operator
+
+        Returns:
+            list[Operator]: ancestors of the given operators
+        """
+
+        ancestors = {i for ind in indexes for i in rx.ancestors(self._graph, ind)}
+        if sort:
+            ancestors = sorted(ancestors)
+        return [self._queue[ind] for ind in ancestors]
+
+    def descendants_of_indexes(self, indexes: Sequence[int], sort=False):
+        """Descendants of a given set of operators.
+
+        Args:
+            indexes (Sequence[int]) : the index into the queue for the operator
+
+        Returns:
+            list[Operator]: descendants of the given operators
+        """
+
+        ancestors = {i for ind in indexes for i in rx.descendants(self._graph, ind)}
+        if sort:
+            ancestors = sorted(ancestors)
+        return [self._queue[ind] for ind in ancestors]
+
     def descendants(self, ops, sort=False):
         """Descendants of a given set of operators.
 
@@ -289,7 +321,8 @@ class CircuitGraph:
             )
         if any(len(self._inds_for_objs[WrappedObj(op)]) > 1 for op in ops):
             raise ValueError(
-                "cannot calculate decendents for an operator that occurs multiple times."
+                "cannot calculate decendents for an operator that occurs multiple times. "
+                "Please use descendants_of_indexes instead."
             )
         descendants = set()
         for op in ops:
@@ -356,25 +389,26 @@ class CircuitGraph:
         """
         # FIXME maybe layering should be greedier, for example [a0 b0 c1 d1] should layer as [a0
         # c1], [b0, d1] and not [a0], [b0 c1], [d1] keep track of the current layer
-        current = Layer([], [])
+        current = Layer([], [], [])
         layers = [current]
 
         for idx, info in enumerate(self.par_info):
-            if idx in self.trainable_params:
+            if self.trainable_params and idx in self.trainable_params:
                 op = info["op"]
 
                 # get all predecessor ops of the op
-                sub = self.ancestors((op,))
+                sub = self.ancestors_of_indexes((info["op_idx"],))
 
                 # check if any of the dependents are in the
                 # currently assembled layer
                 if any(o1 is o2 for o1 in current.ops for o2 in sub):
                     # operator depends on current layer, start a new layer
-                    current = Layer([], [])
+                    current = Layer([], [], [])
                     layers.append(current)
 
                 # store the parameters and ops indices for the layer
                 current.ops.append(op)
+                current.ops_inds.append(info["op_idx"])
                 current.param_inds.append(idx)
 
         return layers
@@ -386,9 +420,9 @@ class CircuitGraph:
             Iterable[LayerData]: layers with extra metadata
         """
         # iterate through each layer
-        for ops, param_inds in self.parametrized_layers:
-            pre_queue = self.ancestors_in_order(ops)
-            post_queue = self.descendants_in_order(ops)
+        for ops, param_inds, indexes in self.parametrized_layers:
+            pre_queue = self.ancestors_of_indexes(indexes, sort=True)
+            post_queue = self.descendants_of_indexes(indexes, sort=True)
             yield LayerData(pre_queue, ops, tuple(param_inds), post_queue)
 
     def update_node(self, old, new):


Test patch:
diff --git a/tests/circuit_graph/test_circuit_graph.py b/tests/circuit_graph/test_circuit_graph.py
index d28dc2341..28197c0ad 100644
--- a/tests/circuit_graph/test_circuit_graph.py
+++ b/tests/circuit_graph/test_circuit_graph.py
@@ -163,19 +163,23 @@ class TestCircuitGraph:
         queue = ops + obs
 
         ancestors = circuit.ancestors([queue[6]])
-        assert len(ancestors) == 3
+        ancestors_index = circuit.ancestors_of_indexes([6])
+        assert len(ancestors) == len(ancestors_index) == 3
         for o_idx in (0, 1, 3):
             assert queue[o_idx] in ancestors
+            assert queue[o_idx] in ancestors_index
 
         descendants = circuit.descendants([queue[6]])
+        descendants_index = circuit.descendants_of_indexes([6])
         assert descendants == [queue[8]]
+        assert descendants_index == [queue[8]]
 
     def test_ancestors_and_descendents_repeated_op(self):
         """Test ancestors and descendents raises a ValueError is the requested operation occurs more than once."""
 
         op = qml.X(0)
         ops = [op, qml.Y(0), op, qml.Z(0), op]
-        graph = CircuitGraph(ops, [], [0, 1, 2])
+        graph = CircuitGraph(ops, [], qml.wires.Wires([0, 1, 2]))
 
         with pytest.raises(ValueError, match=r"operator that occurs multiple times."):
             graph.ancestors([op])
@@ -268,6 +272,27 @@ class TestCircuitGraph:
         assert layers[2].ops == [ops[x] for x in [5, 6]]
         assert layers[2].param_inds == [6, 7]
 
+    def test_iterate_layers_repeat_op(self):
+        """Test iterate_parametrized_layers can work when the operation is repeated."""
+        op = qml.RX(0.5, 0)
+        par_info = [{"op": op, "op_idx": 0, "p_idx": 0}, {"op": op, "op_idx": 2, "p_idx": 0}]
+        graph = qml.CircuitGraph(
+            [op, qml.X(0), op], [], wires=op.wires, trainable_params={0, 1}, par_info=par_info
+        )
+        layers = list(graph.iterate_parametrized_layers())
+
+        assert len(layers) == 2
+
+        assert layers[0].pre_ops == []
+        assert layers[0].ops == [op]
+        assert layers[0].param_inds == (0,)
+        assert layers[0].post_ops == [qml.X(0), op]
+
+        assert layers[1].ops == [op]
+        assert layers[1].param_inds == (1,)
+        assert layers[1].pre_ops == [op, qml.X(0)]
+        assert layers[1].post_ops == []
+
     @pytest.mark.parametrize("wires", [["a", "q1", 3]])
     def test_iterate_layers(self, parametrized_circuit_gaussian, wires):
         """A test of the different layers, their successors and ancestors using a simple circuit"""
diff --git a/tests/gradients/core/test_metric_tensor.py b/tests/gradients/core/test_metric_tensor.py
index 7ea63e56e..12d2c471f 100644
--- a/tests/gradients/core/test_metric_tensor.py
+++ b/tests/gradients/core/test_metric_tensor.py
@@ -1650,3 +1650,16 @@ def test_get_aux_wire_with_unavailable_aux():
     device_wires = qml.wires.Wires([0, "one"])
     with pytest.raises(qml.wires.WireError, match="The requested auxiliary wire does not exist"):
         _get_aux_wire("two", tape, device_wires)
+
+
+def test_metric_tensor_repeated_parametrized_op():
+    """Test that metric tensor works when an operator is repeated."""
+    op = qml.RX(0.5, 0)
+    tape1 = qml.tape.QuantumScript([op, op], [qml.expval(qml.Z(0))])
+    tape2 = qml.tape.QuantumScript([op, qml.RX(0.5, 0)], [qml.expval(qml.Z(0))])
+
+    batch1, _ = qml.metric_tensor(tape1)
+    batch2, _ = qml.metric_tensor(tape2)
+
+    for t1, t2 in zip(batch1, batch2, strict=True):
+        qml.assert_equal(t1, t2)