Repository: PennyLaneAI/pennylane
Pull request #5801 (2024-06-05 10:16:49)

Added repetitions to UCCSD template
**Context:**
As discussed in #5741, the UCCSD template should accept an optional argument `reps` which defines the number of times the UCCSD template is repeated

**Description of the Change:**
`qml.UCCSD` now accepts an additional optional argument `reps` which is an integer and defines the number of times the UCCSD template is repeated. This can improve the accuracy of the UCCSD template by reducing the Trotter-error but results in deeper circuits. The provided weights should have the shape `(reps, len(s_wires) + len(d_wires),)`. If `reps=1`, which is the default value, the weights can also have the shape `(len(s_wires) + len(d_wires),)`. Therefore, already existing code using the UCCSD template does not break.

**Benefits:**
 The change can improve the accuracy of the UCCSD template by reducing the Trotter-error.

**Possible Drawbacks:**
 The change results in deeper circuits depending on the value of `reps`.

**Related GitHub Issues:**
#5741 


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5801?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `86.66667%` with `2 lines` in your changes missing coverage. Please review.
> Project coverage is 99.66%. Comparing base [(`3f0dec4`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/3f0dec439dc841d7ba68dd1422957d023d737681?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`ffdb477`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/ffdb4774e4faa2f144de704d80cf810f5c13f1f9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5801?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/templates/subroutines/uccsd.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5801?src=pr&el=tree&filepath=pennylane%2Ftemplates%2Fsubroutines%2Fuccsd.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdWJyb3V0aW5lcy91Y2NzZC5weQ==) | 86.66% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5801?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5801      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.02%     
==========================================
  Files         414      414              
  Lines       39463    39175     -288     
==========================================
- Hits        39336    39044     -292     
- Misses        127      131       +4     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5801?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

trbromley: Thanks a lot @S-Erik! I've tagged someone from our team to take a look.
S-Erik: I added a few more test cases in `test_uccsd.py`. Please take a look if this is sufficient or if I should add more.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 60b40c8dd..c148e6c83 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -8,6 +8,11 @@
   
 <h3>Improvements 🛠</h3>
 
+* `qml.UCCSD` now accepts an additional optional argument, `n_repeats`, which defines the number of
+  times the UCCSD template is repeated. This can improve the accuracy of the template by reducing
+  the Trotter error but would result in deeper circuits.
+  [(#5801)](https://github.com/PennyLaneAI/pennylane/pull/5801)
+
 * `QuantumScript.hash` is now cached, leading to performance improvements.
   [(#5919)](https://github.com/PennyLaneAI/pennylane/pull/5919)
 
@@ -24,4 +29,5 @@
 This release contains contributions from (in alphabetical order):
 
 Yushao Chen,
-Christina Lee,
\ No newline at end of file
+Christina Lee,
+Erik Schultheis.
diff --git a/pennylane/templates/subroutines/uccsd.py b/pennylane/templates/subroutines/uccsd.py
index bfaaf1307..af5135f7f 100644
--- a/pennylane/templates/subroutines/uccsd.py
+++ b/pennylane/templates/subroutines/uccsd.py
@@ -63,13 +63,16 @@ class UCCSD(Operation):
         \{\mathrm{H.c.}\}) \Big\}.
 
     Args:
-        weights (tensor_like): Size ``(len(s_wires) + len(d_wires),)`` tensor containing the parameters
+        weights (tensor_like): Size ``(n_repeats, len(s_wires) + len(d_wires),)`` tensor containing the
+            parameters (see usage details below) :math:`\theta_{pr}` and :math:`\theta_{pqrs}` entering
+            the Z rotation in :func:`~.FermionicSingleExcitation` and :func:`~.FermionicDoubleExcitation`.
+            These parameters are the coupled-cluster amplitudes that need to be optimized for each
+            single and double excitation generated with the :func:`~.excitations` function.
+            If the size of the given tensor is ``(len(s_wires) + len(d_wires),)``, it is assumed that ``n_repeats == 1``.
             :math:`\theta_{pr}` and :math:`\theta_{pqrs}` entering the Z rotation in
-            :func:`~.FermionicSingleExcitation`
-            and
-            :func:`~.FermionicDoubleExcitation`. These parameters are the coupled-cluster
-            amplitudes that need to be optimized for each single and double excitation generated
-            with the :func:`~.excitations` function.
+            :func:`~.FermionicSingleExcitation` and :func:`~.FermionicDoubleExcitation`.
+            These parameters are the coupled-cluster amplitudes that need to be optimized for each
+            single and double excitation generated with the :func:`~.excitations` function.
         wires (Iterable): wires that the template acts on
         s_wires (Sequence[Sequence]): Sequence of lists containing the wires ``[r,...,p]``
             resulting from the single excitation
@@ -87,6 +90,7 @@ class UCCSD(Operation):
             and unoccupied orbitals in the intervals ``[s, r]`` and ``[q, p]``, respectively.
         init_state (array[int]): Length ``len(wires)`` occupation-number vector representing the
             HF state. ``init_state`` is used to initialize the wires.
+        n_repeats (int): Number of times the UCCSD unitary is repeated. The default value is ``1``.
 
     .. details::
         :title: Usage Details
@@ -99,8 +103,10 @@ class UCCSD(Operation):
         #. The single and double excitations can be generated with the function
            :func:`~.excitations`. See example below.
 
-        #. The vector of parameters ``weights`` is a one-dimensional array of size
-           ``len(s_wires)+len(d_wires)``
+        #. The vector of parameters ``weights`` is a two-dimensional array of size
+           ``(n_repeats, len(s_wires)+len(d_wires))``.
+        #. If ``n_repeats=1``, then ``weights`` can also be a one-dimensional array of size
+           ``len(s_wires)+len(d_wires)``.
 
 
         An example of how to use this template is shown below:
@@ -171,7 +177,9 @@ class UCCSD(Operation):
     num_wires = AnyWires
     grad_method = None
 
-    def __init__(self, weights, wires, s_wires=None, d_wires=None, init_state=None, id=None):
+    def __init__(
+        self, weights, wires, s_wires=None, d_wires=None, init_state=None, n_repeats=1, id=None
+    ):
         if (not s_wires) and (not d_wires):
             raise ValueError(
                 f"s_wires and d_wires lists can not be both empty; got ph={s_wires}, pphh={d_wires}"
@@ -183,10 +191,20 @@ class UCCSD(Operation):
                     f"expected entries of d_wires to be of size 2; got {d_wires_} of length {len(d_wires_)}"
                 )
 
+        if n_repeats < 1:
+            raise ValueError(f"Requires n_repeats to be at least 1; got {n_repeats}.")
+
         shape = qml.math.shape(weights)
-        if shape != (len(s_wires) + len(d_wires),):
+
+        expected_shape = (len(s_wires) + len(d_wires),)
+        if len(shape) == 1 and (n_repeats != 1 or shape != expected_shape):
+            raise ValueError(
+                f"For one-dimensional weights tensor, the shape must be {expected_shape}, and n_repeats should be 1; "
+                f"got {shape} and {n_repeats}, respectively."
+            )
+        if len(shape) != 1 and shape != (n_repeats,) + expected_shape:
             raise ValueError(
-                f"Weights tensor must be of shape {(len(s_wires) + len(d_wires),)}; got {shape}."
+                f"Weights tensor must be of shape {(n_repeats,) + expected_shape}; got {shape}."
             )
 
         init_state = qml.math.toarray(init_state)
@@ -198,6 +216,7 @@ class UCCSD(Operation):
             "init_state": tuple(init_state),
             "s_wires": tuple(tuple(w) for w in s_wires),
             "d_wires": tuple(tuple(tuple(w) for w in dw) for dw in d_wires),
+            "n_repeats": n_repeats,
         }
 
         super().__init__(weights, wires=wires, id=id)
@@ -220,7 +239,7 @@ class UCCSD(Operation):
 
     @staticmethod
     def compute_decomposition(
-        weights, wires, s_wires, d_wires, init_state
+        weights, wires, s_wires, d_wires, init_state, n_repeats
     ):  # pylint: disable=arguments-differ
         r"""Representation of the operator as a product of other operators.
 
@@ -231,7 +250,8 @@ class UCCSD(Operation):
         .. seealso:: :meth:`~.UCCSD.decomposition`.
 
         Args:
-            weights (tensor_like): Size ``(len(s_wires) + len(d_wires),)`` tensor containing the parameters
+            weights (tensor_like): Size ``(len(s_wires) + len(d_wires),)`` or ``(n_repeats, len(s_wires) + len(d_wires),)``,
+                depending on ``n_repeats``, tensor containing the parameters
                 entering the Z rotation in :func:`~.FermionicSingleExcitation` and :func:`~.FermionicDoubleExcitation`.
             wires (Any or Iterable[Any]): wires that the operator acts on
             s_wires (Sequence[Sequence]): Sequence of lists containing the wires ``[r,...,p]``
@@ -240,6 +260,7 @@ class UCCSD(Operation):
                 specify the indices ``[s, ...,r]`` and ``[q,..., p]`` defining the double excitation.
             init_state (array[int]): Length ``len(wires)`` occupation-number vector representing the
                 HF state. ``init_state`` is used to initialize the wires.
+            n_repeats (int): Number of times the UCCSD unitary is repeated.
 
         Returns:
             list[.Operator]: decomposition of the operator
@@ -248,12 +269,18 @@ class UCCSD(Operation):
 
         op_list.append(BasisState(init_state, wires=wires))
 
-        for i, (w1, w2) in enumerate(d_wires):
-            op_list.append(
-                qml.FermionicDoubleExcitation(weights[len(s_wires) + i], wires1=w1, wires2=w2)
-            )
+        if n_repeats == 1 and len(qml.math.shape(weights)) == 1:
+            weights = qml.math.expand_dims(weights, 0)
+
+        for layer in range(n_repeats):
+            for i, (w1, w2) in enumerate(d_wires):
+                op_list.append(
+                    qml.FermionicDoubleExcitation(
+                        weights[layer][len(s_wires) + i], wires1=w1, wires2=w2
+                    )
+                )
 
-        for j, s_wires_ in enumerate(s_wires):
-            op_list.append(qml.FermionicSingleExcitation(weights[j], wires=s_wires_))
+            for j, s_wires_ in enumerate(s_wires):
+                op_list.append(qml.FermionicSingleExcitation(weights[layer][j], wires=s_wires_))
 
         return op_list


Test patch:
diff --git a/tests/templates/test_subroutines/test_uccsd.py b/tests/templates/test_subroutines/test_uccsd.py
index da6a3f36f..35d86f03d 100644
--- a/tests/templates/test_subroutines/test_uccsd.py
+++ b/tests/templates/test_subroutines/test_uccsd.py
@@ -27,6 +27,7 @@ test_data_decomposition = [
         [[0, 1, 2]],
         [],
         np.array([3.815]),
+        1,
         [
             [0, qml.BasisState, [0, 1, 2, 3, 4, 5], [np.array([1, 1, 0, 0, 0, 0])]],
             [1, qml.RX, [0], [-np.pi / 2]],
@@ -38,6 +39,7 @@ test_data_decomposition = [
         [[0, 1, 2], [1, 2, 3]],
         [],
         np.array([3.815, 4.866]),
+        1,
         [
             [2, qml.Hadamard, [2], []],
             [8, qml.RX, [0], [np.pi / 2]],
@@ -51,6 +53,7 @@ test_data_decomposition = [
         [],
         [[[0, 1], [2, 3, 4, 5]]],
         np.array([3.815]),
+        1,
         [
             [3, qml.RX, [2], [-np.pi / 2]],
             [29, qml.RZ, [5], [0.476875]],
@@ -64,6 +67,7 @@ test_data_decomposition = [
         [],
         [[[0, 1], [2, 3]], [[0, 1], [4, 5]]],
         np.array([3.815, 4.866]),
+        1,
         [
             [4, qml.Hadamard, [3], []],
             [16, qml.RX, [0], [-np.pi / 2]],
@@ -80,6 +84,7 @@ test_data_decomposition = [
         [[0, 1, 2, 3, 4], [1, 2, 3]],
         [[[0, 1], [2, 3]], [[0, 1], [4, 5]]],
         np.array([3.815, 4.866, 1.019, 0.639]),
+        1,
         [
             [16, qml.RX, [0], [-np.pi / 2]],
             [47, qml.Hadamard, [1], []],
@@ -92,11 +97,38 @@ test_data_decomposition = [
             [235, qml.CNOT, [1, 4], []],
         ],
     ),
+    (
+        [[0, 1, 2]],
+        [],
+        np.array([[3.815]]),
+        1,
+        [
+            [0, qml.BasisState, [0, 1, 2, 3, 4, 5], [np.array([1, 1, 0, 0, 0, 0])]],
+            [1, qml.RX, [0], [-np.pi / 2]],
+            [5, qml.RZ, [2], [1.9075]],
+            [6, qml.CNOT, [1, 2], []],
+        ],
+    ),
+    (
+        [[0, 1, 2]],
+        [],
+        np.array([[3.815], [1.019]]),
+        2,
+        [
+            [0, qml.BasisState, [0, 1, 2, 3, 4, 5], [np.array([1, 1, 0, 0, 0, 0])]],
+            [1, qml.RX, [0], [-np.pi / 2]],
+            [5, qml.RZ, [2], [1.9075]],
+            [6, qml.CNOT, [1, 2], []],
+            [19, qml.RX, [0], [-np.pi / 2]],
+            [23, qml.RZ, [2], [0.5095]],
+            [24, qml.CNOT, [1, 2], []],
+        ],
+    ),
 ]
 
 
-@pytest.mark.parametrize("s_wires, d_wires, weights, _", test_data_decomposition)
-def test_standard_validity(s_wires, d_wires, weights, _):
+@pytest.mark.parametrize("s_wires, d_wires, weights, n_repeats, _", test_data_decomposition)
+def test_standard_validity(s_wires, d_wires, weights, n_repeats, _):
     """Test standard validity criteria using assert_valid."""
     cnots = 0
     for s_wires_ in s_wires:
@@ -104,25 +136,37 @@ def test_standard_validity(s_wires, d_wires, weights, _):
 
     for d_wires_ in d_wires:
         cnots += 16 * (len(d_wires_[0]) - 1 + len(d_wires_[1]) - 1 + 1)
+
+    cnots *= n_repeats
+
     N = 6
     wires = range(N)
 
     ref_state = np.array([1, 1, 0, 0, 0, 0])
 
-    op = qml.UCCSD(weights, wires, s_wires=s_wires, d_wires=d_wires, init_state=ref_state)
+    op = qml.UCCSD(
+        weights,
+        wires,
+        s_wires=s_wires,
+        d_wires=d_wires,
+        init_state=ref_state,
+        n_repeats=n_repeats,
+    )
     qml.ops.functions.assert_valid(op)
 
 
 class TestDecomposition:
     """Tests that the template defines the correct decomposition."""
 
-    @pytest.mark.parametrize("s_wires, d_wires, weights, ref_gates", test_data_decomposition)
-    def test_uccsd_operations(self, s_wires, d_wires, weights, ref_gates):
+    @pytest.mark.parametrize(
+        "s_wires, d_wires, weights, n_repeats, ref_gates", test_data_decomposition
+    )
+    def test_uccsd_operations(self, s_wires, d_wires, weights, n_repeats, ref_gates):
         """Test the correctness of the UCCSD template including the gate count
         and order, the wires the operation acts on and the correct use of parameters
         in the circuit."""
 
-        sqg = 10 * len(s_wires) + 72 * len(d_wires)
+        sqg = (10 * len(s_wires) + 72 * len(d_wires)) * n_repeats
 
         cnots = 0
         for s_wires_ in s_wires:
@@ -130,12 +174,22 @@ class TestDecomposition:
 
         for d_wires_ in d_wires:
             cnots += 16 * (len(d_wires_[0]) - 1 + len(d_wires_[1]) - 1 + 1)
+
+        cnots *= n_repeats
+
         N = 6
         wires = range(N)
 
         ref_state = np.array([1, 1, 0, 0, 0, 0])
 
-        op = qml.UCCSD(weights, wires, s_wires=s_wires, d_wires=d_wires, init_state=ref_state)
+        op = qml.UCCSD(
+            weights,
+            wires,
+            s_wires=s_wires,
+            d_wires=d_wires,
+            init_state=ref_state,
+            n_repeats=n_repeats,
+        )
         raw_queue = op.expand().operations
 
         # hack to avoid updating the test data:
@@ -205,13 +259,14 @@ class TestInputs:
     """Test inputs and pre-processing."""
 
     @pytest.mark.parametrize(
-        ("weights", "s_wires", "d_wires", "init_state", "msg_match"),
+        ("weights", "s_wires", "d_wires", "init_state", "n_repeats", "msg_match"),
         [
             (
                 np.array([-2.8]),
                 [[0, 1, 2]],
                 [],
                 np.array([1.2, 1, 0, 0]),
+                1,
                 "Elements of 'init_state' must be integers",
             ),
             (
@@ -219,6 +274,7 @@ class TestInputs:
                 [],
                 [],
                 np.array([1, 1, 0, 0]),
+                1,
                 "s_wires and d_wires lists can not be both empty",
             ),
             (
@@ -226,6 +282,7 @@ class TestInputs:
                 [],
                 [[[0, 1, 2, 3]]],
                 np.array([1, 1, 0, 0]),
+                1,
                 "expected entries of d_wires to be of size 2",
             ),
             (
@@ -233,6 +290,7 @@ class TestInputs:
                 [[0, 2]],
                 [],
                 np.array([1, 1, 0, 0, 0]),
+                1,
                 "Basis states must be of length 4",
             ),
             (
@@ -240,25 +298,60 @@ class TestInputs:
                 [[0, 1, 2]],
                 [],
                 np.array([1, 1, 0, 0]),
-                "Weights tensor must be of",
+                1,
+                "For one-dimensional weights tensor",
             ),
             (
                 np.array([-2.8, 1.6]),
                 [],
                 [[[0, 1], [2, 3]]],
                 np.array([1, 1, 0, 0]),
-                "Weights tensor must be of",
+                1,
+                "For one-dimensional weights tensor",
             ),
             (
                 np.array([-2.8, 1.6]),
                 [[0, 1, 2], [1, 2, 3]],
                 [[[0, 1], [2, 3]]],
                 np.array([1, 1, 0, 0]),
+                1,
+                "For one-dimensional weights tensor",
+            ),
+            (
+                np.array([-2.8]),
+                [[0, 1, 2]],
+                [],
+                np.array([1, 1, 0, 0]),
+                0,
+                "Requires n_repeats to be at least 1",
+            ),
+            (
+                np.array([-2.8]),
+                [[0, 1, 2]],
+                [],
+                np.array([1, 1, 0, 0]),
+                -1,
+                "Requires n_repeats to be at least 1",
+            ),
+            (
+                np.array([-2.8]),
+                [[0, 1, 2]],
+                [],
+                np.array([1, 1, 0, 0]),
+                2,
+                "For one-dimensional weights tensor",
+            ),
+            (
+                np.array([[-2.8], [-1.8]]),
+                [[0, 1, 2]],
+                [],
+                np.array([1, 1, 0, 0]),
+                3,
                 "Weights tensor must be of",
             ),
         ],
     )
-    def test_uccsd_xceptions(self, weights, s_wires, d_wires, init_state, msg_match):
+    def test_uccsd_xceptions(self, weights, s_wires, d_wires, init_state, n_repeats, msg_match):
         """Test that UCCSD throws an exception if the parameters have illegal
         shapes, types or values."""
         N = 4
@@ -266,7 +359,12 @@ class TestInputs:
         dev = qml.device("default.qubit", wires=N)
 
         def circuit(
-            weights=weights, wires=wires, s_wires=s_wires, d_wires=d_wires, init_state=init_state
+            weights=weights,
+            wires=wires,
+            s_wires=s_wires,
+            d_wires=d_wires,
+            init_state=init_state,
+            n_repeats=n_repeats,
         ):
             qml.UCCSD(
                 weights=weights,
@@ -274,6 +372,7 @@ class TestInputs:
                 s_wires=s_wires,
                 d_wires=d_wires,
                 init_state=init_state,
+                n_repeats=n_repeats,
             )
             return qml.expval(qml.PauliZ(0))
 
@@ -286,6 +385,7 @@ class TestInputs:
                 s_wires=s_wires,
                 d_wires=d_wires,
                 init_state=init_state,
+                n_repeats=n_repeats,
             )
 
     def test_id(self):
@@ -301,21 +401,27 @@ class TestInputs:
         assert template.id == "a"
 
 
-def circuit_template(weights):
+def circuit_template(weights, n_repeats=1):
     qml.UCCSD(
         weights,
         wires=range(4),
         s_wires=[[0, 1]],
         d_wires=[[[0, 1], [2, 3]]],
         init_state=np.array([1, 0, 0, 0]),
+        n_repeats=n_repeats,
     )
     return qml.expval(qml.PauliZ(0))
 
 
-def circuit_decomposed(weights):
+def circuit_decomposed(weights, n_repeats=1):
     qml.BasisState(np.array([1, 0, 0, 0]), wires=range(4))
-    qml.FermionicDoubleExcitation(weights[1], wires1=[0, 1], wires2=[2, 3])
-    qml.FermionicSingleExcitation(weights[0], wires=[0, 1])
+
+    if n_repeats == 1 and len(qml.math.shape(weights)) == 1:
+        weights = qml.math.expand_dims(weights, 0)
+
+    for layer in range(n_repeats):
+        qml.FermionicDoubleExcitation(weights[layer][1], wires1=[0, 1], wires2=[2, 3])
+        qml.FermionicSingleExcitation(weights[layer][0], wires=[0, 1])
     return qml.expval(qml.PauliZ(0))
 
 
@@ -337,6 +443,12 @@ class TestInterfaces:
         res2 = circuit2(weights)
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
 
+        # Test with n_repeats=2
+        weights = [list(range(2)), list(range(2))]
+        res = circuit(weights, n_repeats=2)
+        res2 = circuit2(weights, n_repeats=2)
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
     @pytest.mark.autograd
     def test_autograd(self, tol):
         """Tests the autograd interface."""
@@ -360,6 +472,28 @@ class TestInterfaces:
 
         assert np.allclose(grads, grads2, atol=tol, rtol=0)
 
+        # Test with n_repeats=2
+        weights = pnp.array(
+            np.random.random(
+                size=(
+                    2,
+                    2,
+                )
+            ),
+            requires_grad=True,
+        )
+        res = circuit(weights, n_repeats=2)
+        res2 = circuit2(weights, n_repeats=2)
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+        grad_fn = qml.grad(circuit)
+        grads = grad_fn(weights, n_repeats=2)
+
+        grad_fn2 = qml.grad(circuit2)
+        grads2 = grad_fn2(weights, n_repeats=2)
+
+        assert np.allclose(grads, grads2, atol=tol, rtol=0)
+
     @pytest.mark.jax
     @pytest.mark.slow
     def test_jax(self, tol):
@@ -386,6 +520,27 @@ class TestInterfaces:
 
         assert np.allclose(grads[0], grads2[0], atol=tol, rtol=0)
 
+        # Test with n_repeats=2
+        weights = jnp.array(
+            np.random.random(
+                size=(
+                    2,
+                    2,
+                )
+            )
+        )
+        res = circuit(weights, n_repeats=2)
+        res2 = circuit2(weights, n_repeats=2)
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+        grad_fn = jax.grad(circuit)
+        grads = grad_fn(weights, n_repeats=2)
+
+        grad_fn2 = jax.grad(circuit2)
+        grads2 = grad_fn2(weights, n_repeats=2)
+
+        assert np.allclose(grads[0], grads2[0], atol=tol, rtol=0)
+
     @pytest.mark.tf
     def test_tf(self, tol):
         """Tests the tf interface."""
@@ -412,6 +567,29 @@ class TestInterfaces:
 
         assert np.allclose(grads[0], grads2[0], atol=tol, rtol=0)
 
+        # Test with n_repeats=2
+        weights = tf.Variable(
+            np.random.random(
+                size=(
+                    2,
+                    2,
+                )
+            )
+        )
+        res = circuit(weights, n_repeats=2)
+        res2 = circuit2(weights, n_repeats=2)
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+        with tf.GradientTape() as tape:
+            res = circuit(weights, n_repeats=2)
+        grads = tape.gradient(res, [weights])
+
+        with tf.GradientTape() as tape2:
+            res2 = circuit2(weights, n_repeats=2)
+        grads2 = tape2.gradient(res2, [weights])
+
+        assert np.allclose(grads[0], grads2[0], atol=tol, rtol=0)
+
     @pytest.mark.torch
     def test_torch(self, tol):
         """Tests the torch interface."""
@@ -438,3 +616,30 @@ class TestInterfaces:
         grads2 = [weights.grad]
 
         assert np.allclose(grads[0], grads2[0], atol=tol, rtol=0)
+
+        # Test with n_repeats=2
+        weights = torch.tensor(
+            np.random.random(
+                size=(
+                    2,
+                    2,
+                )
+            ),
+            requires_grad=True,
+        )
+        circuit = qml.QNode(circuit_template, dev)
+        circuit2 = qml.QNode(circuit_decomposed, dev)
+
+        res = circuit(weights, n_repeats=2)
+        res2 = circuit2(weights, n_repeats=2)
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+        res = circuit(weights, n_repeats=2)
+        res.backward()
+        grads = [weights.grad]
+
+        res2 = circuit2(weights, n_repeats=2)
+        res2.backward()
+        grads2 = [weights.grad]
+
+        assert np.allclose(grads[0], grads2[0], atol=tol, rtol=0)