Repository: PennyLaneAI/pennylane
Pull request #5523 (2024-04-16 10:52:50)

Plxpr integrates well with all templates
**To do**
- [x] Actually implement the custom binding functions
- [x] Write the custom binding tests
- [x] Write a test that asserts that all modified templates as determined automatically have their own test.

**Context:**
Introducing the plxpr capturing mechanism requires us to pay attention to operations that have special call signatures, such as templates.

**Description of the Change:**
This PR adds a test suite for templates.
For templates with special initialization signatures, a custom `_primitive_bind_call` method is added, together with dedicated tests for these custom methods.

**Benefits:**
Compatibility of Plxpr with the code stack

**Possible Drawbacks:**
More than 1 second of testing time added.
A lot of test cases
Future template additions _will_ have to write additional tests in `tests/capture/test_templates.py` in the current test suite design.

**Related GitHub Issues:**


[sc-61425]

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
dime10: > @dime10 Yes you are right about the wires binding, I think. I can try and refine this PR to respect those special settings of split wires. Alternatively, we could go through with this PR and consider it an upgrade to allow dynamic wires in those templates 🙃

That's fine by me, we can leave it up to a follow-up pr to support dynamic wire sets split across multiple arguments :)
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5523?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 `93.54839%` with `4 lines` in your changes are missing coverage. Please review.
> Project coverage is 99.65%. Comparing base [(`761c455`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/761c455e1be44d7b2590e3ce621a1171ab457487?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`d73e791`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/d73e7918332616d1ec9a87d25cfb5424886f152f?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/5523?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/basis\_rotation.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5523?src=pr&el=tree&filepath=pennylane%2Ftemplates%2Fsubroutines%2Fbasis_rotation.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdWJyb3V0aW5lcy9iYXNpc19yb3RhdGlvbi5weQ==) | 84.61% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5523?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |
| [pennylane/templates/subroutines/hilbert\_schmidt.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5523?src=pr&el=tree&filepath=pennylane%2Ftemplates%2Fsubroutines%2Fhilbert_schmidt.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdWJyb3V0aW5lcy9oaWxiZXJ0X3NjaG1pZHQucHk=) | 83.33% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5523?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |
| [pennylane/templates/subroutines/qubitization.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5523?src=pr&el=tree&filepath=pennylane%2Ftemplates%2Fsubroutines%2Fqubitization.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdWJyb3V0aW5lcy9xdWJpdGl6YXRpb24ucHk=) | 66.66% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5523?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                 @@
##           valid-templates    #5523      +/-   ##
===================================================
- Coverage            99.66%   99.65%   -0.01%     
===================================================
  Files                  416      416              
  Lines                38536    38597      +61     
===================================================
+ Hits                 38408    38465      +57     
- Misses                 128      132       +4     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5523?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).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 323549d81..489fd740d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -102,6 +102,7 @@
 * PennyLane operators and measurements can now automatically be captured as instructions in JAXPR.
   [(#5564)](https://github.com/PennyLaneAI/pennylane/pull/5564)
   [(#5511)](https://github.com/PennyLaneAI/pennylane/pull/5511)
+  [(#5523)](https://github.com/PennyLaneAI/pennylane/pull/5523)
 
 * The `decompose` transform has an `error` kwarg to specify the type of error that should be raised, 
   allowing error types to be more consistent with the context the `decompose` function is used in.
diff --git a/pennylane/capture/__init__.py b/pennylane/capture/__init__.py
index 6cd2dacf4..33edf07db 100644
--- a/pennylane/capture/__init__.py
+++ b/pennylane/capture/__init__.py
@@ -81,7 +81,7 @@ But an operator developer may need to override custom behavior for calling ``cls
 (where ``cls`` indicates the class) if:
 
 * The operator does not accept wires, like :class:`~.SymbolicOp` or :class:`~.CompositeOp`.
-* The operator needs to enforce a data/ metadata distinction, like :class:`~.PauliRot`.
+* The operator needs to enforce a data / metadata distinction, like :class:`~.PauliRot`.
 
 In such cases, the operator developer can override ``cls._primitive_bind_call``, which
 will be called when constructing a new class instance instead of ``type.__call__``.  For example,
diff --git a/pennylane/templates/layers/gate_fabric.py b/pennylane/templates/layers/gate_fabric.py
index 7276dc510..5f3a41247 100644
--- a/pennylane/templates/layers/gate_fabric.py
+++ b/pennylane/templates/layers/gate_fabric.py
@@ -281,7 +281,7 @@ class GateFabric(Operation):
 
         if n_wires < 4:
             raise ValueError(
-                f"This template requires the number of qubits to be greater than four; got 'n_wires' = {n_wires}"
+                f"This template requires the number of qubits to be at least four; got 'n_wires' = {n_wires}"
             )
 
         if n_wires % 2:
diff --git a/pennylane/templates/state_preparations/mottonen.py b/pennylane/templates/state_preparations/mottonen.py
index 9bc5269e5..171a572c2 100644
--- a/pennylane/templates/state_preparations/mottonen.py
+++ b/pennylane/templates/state_preparations/mottonen.py
@@ -308,7 +308,7 @@ class MottonenStatePreparation(Operation):
                 norm = qml.math.sum(qml.math.abs(state) ** 2)
                 if not qml.math.allclose(norm, 1.0, atol=1e-3):
                     raise ValueError(
-                        f"State vectors have to be of norm 1.0, vector {i} has norm {norm}"
+                        f"State vectors have to be of norm 1.0, vector {i} has squared norm {norm}"
                     )
 
         super().__init__(state_vector, wires=wires, id=id)
diff --git a/pennylane/templates/subroutines/all_singles_doubles.py b/pennylane/templates/subroutines/all_singles_doubles.py
index 9350a9c11..97d7c5305 100644
--- a/pennylane/templates/subroutines/all_singles_doubles.py
+++ b/pennylane/templates/subroutines/all_singles_doubles.py
@@ -145,7 +145,7 @@ class AllSinglesDoubles(Operation):
             raise ValueError(f"'weights' tensor must be of shape {exp_shape}; got {weights_shape}.")
 
         if hf_state[0].dtype != np.dtype("int"):
-            raise ValueError(f"Elements of 'hf_state' must be integers; got {hf_state.dtype}")
+            raise ValueError(f"Elements of 'hf_state' must be integers; got {hf_state[0].dtype}")
 
         singles = tuple(tuple(s) for s in singles)
         doubles = tuple(tuple(d) for d in doubles)
diff --git a/pennylane/templates/subroutines/amplitude_amplification.py b/pennylane/templates/subroutines/amplitude_amplification.py
index f81192188..08e050ff7 100644
--- a/pennylane/templates/subroutines/amplitude_amplification.py
+++ b/pennylane/templates/subroutines/amplitude_amplification.py
@@ -107,6 +107,10 @@ class AmplitudeAmplification(Operation):
         metadata = tuple(item for item in self.hyperparameters.items() if item[0] not in ["O", "U"])
         return data, metadata
 
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
     @classmethod
     def _unflatten(cls, data, metadata):
         return cls(*data, **dict(metadata))
diff --git a/pennylane/templates/subroutines/approx_time_evolution.py b/pennylane/templates/subroutines/approx_time_evolution.py
index 92d153c1f..2272e37c0 100644
--- a/pennylane/templates/subroutines/approx_time_evolution.py
+++ b/pennylane/templates/subroutines/approx_time_evolution.py
@@ -124,6 +124,10 @@ class ApproxTimeEvolution(Operation):
         data = (h, self.data[-1])
         return data, (self.hyperparameters["n"],)
 
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
     @classmethod
     def _unflatten(cls, data, metadata):
         return cls(data[0], data[1], n=metadata[0])
diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index 16e4dd07c..346e8b837 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -100,6 +100,15 @@ class BasisRotation(Operation):
     num_wires = AnyWires
     grad_method = None
 
+    @classmethod
+    def _primitive_bind_call(cls, wires, unitary_matrix, check=False, id=None):
+        # pylint: disable=arguments-differ
+        if cls._primitive is None:
+            # guard against this being called when primitive is not defined.
+            return type.__call__(cls, wires, unitary_matrix, check=check, id=id)  # pragma: no cover
+
+        return cls._primitive.bind(*wires, unitary_matrix, check=check, id=id)
+
     def __init__(self, wires, unitary_matrix, check=False, id=None):
         M, N = unitary_matrix.shape
         if M != N:
@@ -176,3 +185,18 @@ class BasisRotation(Operation):
                 op_list.append(qml.PhaseShift(phi, wires=wires[indices[0]]))
 
         return op_list
+
+
+# Program capture needs to unpack and re-pack the wires to support dynamic wires. For
+# BasisRotation, the unconventional argument ordering requires custom def_impl code.
+# See capture module for more information on primitives
+# If None, jax isn't installed so the class never got a primitive.
+if BasisRotation._primitive is not None:  # pylint: disable=protected-access
+
+    @BasisRotation._primitive.def_impl  # pylint: disable=protected-access
+    def _(*args, **kwargs):
+        # If there are more than two args, we are calling with unpacked wires, so that
+        # we have to repack them. This replaces the n_wires logic in the general case.
+        if len(args) != 2:
+            args = (args[:-1], args[-1])
+        return type.__call__(BasisRotation, *args, **kwargs)
diff --git a/pennylane/templates/subroutines/commuting_evolution.py b/pennylane/templates/subroutines/commuting_evolution.py
index 19b983b35..1cb5ed003 100644
--- a/pennylane/templates/subroutines/commuting_evolution.py
+++ b/pennylane/templates/subroutines/commuting_evolution.py
@@ -114,6 +114,10 @@ class CommutingEvolution(Operation):
         data = (self.data[0], h)
         return data, (self.hyperparameters["frequencies"], self.hyperparameters["shifts"])
 
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
     @classmethod
     def _unflatten(cls, data, metadata) -> "CommutingEvolution":
         return cls(data[1], data[0], frequencies=metadata[0], shifts=metadata[1])
diff --git a/pennylane/templates/subroutines/fermionic_double_excitation.py b/pennylane/templates/subroutines/fermionic_double_excitation.py
index 262034398..797189dbe 100644
--- a/pennylane/templates/subroutines/fermionic_double_excitation.py
+++ b/pennylane/templates/subroutines/fermionic_double_excitation.py
@@ -504,6 +504,10 @@ class FermionicDoubleExcitation(Operation):
     def _flatten(self):
         return self.data, (self.hyperparameters["wires1"], self.hyperparameters["wires2"])
 
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
     @classmethod
     def _unflatten(cls, data, metadata) -> "FermionicDoubleExcitation":
         return cls(data[0], wires1=metadata[0], wires2=metadata[1])
diff --git a/pennylane/templates/subroutines/hilbert_schmidt.py b/pennylane/templates/subroutines/hilbert_schmidt.py
index 56987a05b..0837766e0 100644
--- a/pennylane/templates/subroutines/hilbert_schmidt.py
+++ b/pennylane/templates/subroutines/hilbert_schmidt.py
@@ -105,6 +105,12 @@ class HilbertSchmidt(Operation):
         )
         return self.data, metadata
 
+    @classmethod
+    def _primitive_bind_call(cls, *params, v_function, v_wires, u_tape, id=None):
+        # pylint: disable=arguments-differ
+        kwargs = {"v_function": v_function, "v_wires": v_wires, "u_tape": u_tape, "id": id}
+        return cls._primitive.bind(*params, **kwargs)
+
     @classmethod
     def _unflatten(cls, data, metadata):
         return cls(*data, **dict(metadata))
diff --git a/pennylane/templates/subroutines/qdrift.py b/pennylane/templates/subroutines/qdrift.py
index 6b7577a8b..a65152586 100644
--- a/pennylane/templates/subroutines/qdrift.py
+++ b/pennylane/templates/subroutines/qdrift.py
@@ -180,6 +180,10 @@ class QDrift(Operation):
 
     """
 
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
     def _flatten(self):
         h = self.hyperparameters["base"]
         hashable_hyperparameters = tuple(
diff --git a/pennylane/templates/subroutines/qmc.py b/pennylane/templates/subroutines/qmc.py
index 1ddc2a433..50a764bc0 100644
--- a/pennylane/templates/subroutines/qmc.py
+++ b/pennylane/templates/subroutines/qmc.py
@@ -344,6 +344,10 @@ class QuantumMonteCarlo(Operation):
     num_wires = AnyWires
     grad_method = None
 
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
     @classmethod
     def _unflatten(cls, data, metadata):
         new_op = cls.__new__(cls)
diff --git a/pennylane/templates/subroutines/qpe.py b/pennylane/templates/subroutines/qpe.py
index 8daff7e6d..0e2e0af44 100644
--- a/pennylane/templates/subroutines/qpe.py
+++ b/pennylane/templates/subroutines/qpe.py
@@ -150,6 +150,10 @@ class QuantumPhaseEstimation(ErrorOperation):
         metadata = (self.hyperparameters["estimation_wires"],)
         return data, metadata
 
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
     @classmethod
     def _unflatten(cls, data, metadata) -> "QuantumPhaseEstimation":
         return cls(data[0], estimation_wires=metadata[0])
diff --git a/pennylane/templates/subroutines/qsvt.py b/pennylane/templates/subroutines/qsvt.py
index d2e896944..fcefba247 100644
--- a/pennylane/templates/subroutines/qsvt.py
+++ b/pennylane/templates/subroutines/qsvt.py
@@ -268,6 +268,10 @@ class QSVT(Operation):
         data = (self.hyperparameters["UA"], self.hyperparameters["projectors"])
         return data, tuple()
 
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
     @classmethod
     def _unflatten(cls, data, _) -> "QSVT":
         return cls(*data)
diff --git a/pennylane/templates/subroutines/qubitization.py b/pennylane/templates/subroutines/qubitization.py
index 5e19a9e9f..ae4548740 100644
--- a/pennylane/templates/subroutines/qubitization.py
+++ b/pennylane/templates/subroutines/qubitization.py
@@ -92,6 +92,10 @@ class Qubitization(Operation):
         eigenvalue: 0.7
     """
 
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
     def __init__(self, hamiltonian, control, id=None):
         wires = hamiltonian.wires + qml.wires.Wires(control)
 
diff --git a/pennylane/templates/subroutines/reflection.py b/pennylane/templates/subroutines/reflection.py
index 7483f5533..2b37ec885 100644
--- a/pennylane/templates/subroutines/reflection.py
+++ b/pennylane/templates/subroutines/reflection.py
@@ -104,6 +104,10 @@ class Reflection(Operation):
 
     """
 
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
     def _flatten(self):
         data = (self.hyperparameters["base"], self.parameters[0])
         return data, (self.hyperparameters["reflection_wires"],)
diff --git a/pennylane/templates/subroutines/select.py b/pennylane/templates/subroutines/select.py
index 107657c85..357cf0113 100644
--- a/pennylane/templates/subroutines/select.py
+++ b/pennylane/templates/subroutines/select.py
@@ -69,6 +69,10 @@ class Select(Operation):
     def _flatten(self):
         return (self.ops), (self.control)
 
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
     @classmethod
     def _unflatten(cls, data, metadata) -> "Select":
         return cls(data, metadata)
diff --git a/pennylane/templates/tensornetworks/mera.py b/pennylane/templates/tensornetworks/mera.py
index 752aea421..5a4b6a1ee 100644
--- a/pennylane/templates/tensornetworks/mera.py
+++ b/pennylane/templates/tensornetworks/mera.py
@@ -175,6 +175,19 @@ class MERA(Operation):
     def num_params(self):
         return 1
 
+    @classmethod
+    def _primitive_bind_call(
+        cls, wires, n_block_wires, block, n_params_block, template_weights=None, id=None
+    ):  # pylint: disable=arguments-differ
+        return super()._primitive_bind_call(
+            wires=wires,
+            n_block_wires=n_block_wires,
+            block=block,
+            n_params_block=n_params_block,
+            template_weights=template_weights,
+            id=id,
+        )
+
     @classmethod
     def _unflatten(cls, data, metadata):
         new_op = cls.__new__(cls)
diff --git a/pennylane/templates/tensornetworks/mps.py b/pennylane/templates/tensornetworks/mps.py
index 40761f4c9..922c5d531 100644
--- a/pennylane/templates/tensornetworks/mps.py
+++ b/pennylane/templates/tensornetworks/mps.py
@@ -163,6 +163,29 @@ class MPS(Operation):
     num_wires = AnyWires
     par_domain = "A"
 
+    @classmethod
+    def _primitive_bind_call(
+        cls,
+        wires,
+        n_block_wires,
+        block,
+        n_params_block,
+        template_weights=None,
+        offset=None,
+        id=None,
+        **kwargs,
+    ):  # pylint: disable=arguments-differ
+        return super()._primitive_bind_call(
+            wires=wires,
+            n_block_wires=n_block_wires,
+            block=block,
+            n_params_block=n_params_block,
+            template_weights=template_weights,
+            id=id,
+            offset=offset,
+            **kwargs,
+        )
+
     @classmethod
     def _unflatten(cls, data, metadata):
         new_op = cls.__new__(cls)
diff --git a/pennylane/templates/tensornetworks/ttn.py b/pennylane/templates/tensornetworks/ttn.py
index 8e274f71d..7bd88ce2e 100644
--- a/pennylane/templates/tensornetworks/ttn.py
+++ b/pennylane/templates/tensornetworks/ttn.py
@@ -146,6 +146,19 @@ class TTN(Operation):
     def num_params(self):
         return 1
 
+    @classmethod
+    def _primitive_bind_call(
+        cls, wires, n_block_wires, block, n_params_block, template_weights=None, id=None
+    ):  # pylint: disable=arguments-differ
+        return super()._primitive_bind_call(
+            wires=wires,
+            n_block_wires=n_block_wires,
+            block=block,
+            n_params_block=n_params_block,
+            template_weights=template_weights,
+            id=id,
+        )
+
     @classmethod
     def _unflatten(cls, data, metadata):
         new_op = cls.__new__(cls)


Test patch:
diff --git a/tests/capture/test_templates.py b/tests/capture/test_templates.py
new file mode 100644
index 000000000..fce97bb77
--- /dev/null
+++ b/tests/capture/test_templates.py
@@ -0,0 +1,758 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Integration tests for the capture of PennyLane templates into plxpr.
+"""
+import inspect
+
+# pylint: disable=protected-access
+from typing import Any
+
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.capture.primitives import _get_abstract_operator
+
+jax = pytest.importorskip("jax")
+jnp = jax.numpy
+
+pytestmark = pytest.mark.jax
+
+AbstractOperator = _get_abstract_operator()
+original_op_bind_code = qml.operation.Operator._primitive_bind_call.__code__
+
+
+@pytest.fixture(autouse=True)
+def enable_disable_plxpr():
+    qml.capture.enable()
+    yield
+    qml.capture.disable()
+
+
+unmodified_templates_cases = [
+    (qml.AmplitudeEmbedding, (jnp.array([1.0, 0.0]), 2), {}),
+    (qml.AmplitudeEmbedding, (jnp.eye(4)[2], [2, 3]), {"normalize": False}),
+    (qml.AmplitudeEmbedding, (jnp.array([0.3, 0.1, 0.2]),), {"pad_with": 1.2, "wires": [0, 3]}),
+    (qml.AngleEmbedding, (jnp.array([1.0, 0.0]), [2, 3]), {}),
+    (qml.AngleEmbedding, (jnp.array([0.4]), [0]), {"rotation": "X"}),
+    (qml.AngleEmbedding, (jnp.array([0.3, 0.1, 0.2]),), {"rotation": "Z", "wires": [0, 2, 3]}),
+    (qml.BasisEmbedding, (jnp.array([1, 0]), [2, 3]), {}),
+    (qml.BasisEmbedding, (), {"features": jnp.array([1, 0]), "wires": [2, 3]}),
+    (qml.BasisEmbedding, (6, [0, 5, 2]), {"id": "my_id"}),
+    (qml.BasisEmbedding, (jnp.array([1, 0, 1]),), {"wires": [0, 2, 3]}),
+    (qml.IQPEmbedding, (jnp.array([2.3, 0.1]), [2, 0]), {}),
+    (qml.IQPEmbedding, (jnp.array([0.4, 0.2, 0.1]), [2, 1, 0]), {"pattern": [[2, 0], [1, 0]]}),
+    (qml.IQPEmbedding, (jnp.array([0.4, 0.1]), [0, 10]), {"n_repeats": 3, "pattern": None}),
+    (qml.QAOAEmbedding, (jnp.array([1.0, 0.0]), jnp.ones((3, 3)), [2, 3]), {}),
+    (qml.QAOAEmbedding, (jnp.array([0.4]), jnp.ones((2, 1)), [0]), {"local_field": "X"}),
+    (
+        qml.QAOAEmbedding,
+        (jnp.array([0.3, 0.1, 0.2]), jnp.zeros((2, 6))),
+        {"local_field": "Z", "wires": [0, 2, 3]},
+    ),
+    (qml.BasicEntanglerLayers, (jnp.ones((5, 2)), [2, 3]), {}),
+    (qml.BasicEntanglerLayers, (jnp.ones((2, 1)), [0]), {"rotation": "X", "id": "my_id"}),
+    (
+        qml.BasicEntanglerLayers,
+        (jnp.array([[0.3, 0.1, 0.2]]),),
+        {"rotation": "Z", "wires": [0, 2, 3]},
+    ),
+    # Need to fix GateFabric positional args: Currently have to pass init_state as kwarg if we want to pass wires as kwarg
+    # https://github.com/PennyLaneAI/pennylane/issues/5521
+    (qml.GateFabric, (jnp.ones((3, 1, 2)), [2, 3, 0, 1]), {"init_state": [0, 1, 1, 0]}),
+    (
+        qml.GateFabric,
+        (jnp.zeros((2, 3, 2)),),
+        {"include_pi": False, "wires": list(range(8)), "init_state": jnp.ones(8)},
+    ),
+    # (qml.GateFabric, (jnp.zeros((2, 3, 2)), jnp.ones(8)), {"include_pi": False, "wires": list(range(8))}), # Can't even init
+    # (qml.GateFabric, (jnp.ones((5, 2, 2)), list(range(6)), jnp.array([0, 0, 1, 1, 0, 1])), {"include_pi": True, "id": "my_id"}), # Can't trace
+    # https://github.com/PennyLaneAI/pennylane/issues/5522
+    # (qml.ParticleConservingU1, (jnp.ones((3, 1, 2)), [2, 3]), {}),
+    (qml.ParticleConservingU1, (jnp.ones((3, 1, 2)), [2, 3]), {"init_state": [0, 1]}),
+    (
+        qml.ParticleConservingU1,
+        (jnp.zeros((5, 3, 2)),),
+        {"wires": [0, 1, 2, 3], "init_state": jnp.ones(4)},
+    ),
+    # https://github.com/PennyLaneAI/pennylane/issues/5522
+    # (qml.ParticleConservingU2, (jnp.ones((3, 3)), [2, 3]), {}),
+    (qml.ParticleConservingU2, (jnp.ones((3, 3)), [2, 3]), {"init_state": [0, 1]}),
+    (
+        qml.ParticleConservingU2,
+        (jnp.zeros((5, 7)),),
+        {"wires": [0, 1, 2, 3], "init_state": jnp.ones(4)},
+    ),
+    (qml.RandomLayers, (jnp.ones((3, 3)), [2, 3]), {}),
+    (qml.RandomLayers, (jnp.ones((3, 3)),), {"wires": [3, 2, 1], "ratio_imprim": 0.5}),
+    (qml.RandomLayers, (), {"weights": jnp.ones((3, 3)), "wires": [3, 2, 1]}),
+    (qml.RandomLayers, (jnp.ones((3, 3)),), {"wires": [3, 2, 1], "rotations": (qml.RX, qml.RZ)}),
+    (qml.RandomLayers, (jnp.ones((3, 3)), [0, 1]), {"rotations": (qml.RX, qml.RZ), "seed": 41}),
+    (qml.SimplifiedTwoDesign, (jnp.ones(2), jnp.zeros((3, 1, 2)), [2, 3]), {}),
+    (qml.SimplifiedTwoDesign, (jnp.ones(3), jnp.zeros((3, 2, 2))), {"wires": [0, 1, 2]}),
+    (qml.SimplifiedTwoDesign, (jnp.ones(2),), {"weights": jnp.zeros((3, 1, 2)), "wires": [0, 2]}),
+    (
+        qml.SimplifiedTwoDesign,
+        (),
+        {"initial_layer_weights": jnp.ones(2), "weights": jnp.zeros((3, 1, 2)), "wires": [0, 2]},
+    ),
+    (qml.StronglyEntanglingLayers, (jnp.ones((3, 2, 3)), [2, 3]), {"ranges": [1, 1, 1]}),
+    (
+        qml.StronglyEntanglingLayers,
+        (jnp.ones((1, 3, 3)),),
+        {"wires": [3, 2, 1], "imprimitive": qml.CZ},
+    ),
+    (qml.StronglyEntanglingLayers, (), {"weights": jnp.ones((3, 3, 3)), "wires": [3, 2, 1]}),
+    (qml.ArbitraryStatePreparation, (jnp.ones(6), [2, 3]), {}),
+    (qml.ArbitraryStatePreparation, (jnp.zeros(14),), {"wires": [3, 2, 0]}),
+    (qml.ArbitraryStatePreparation, (), {"weights": jnp.ones(2), "wires": [1]}),
+    (qml.BasisStatePreparation, (jnp.array([0, 1]), [2, 3]), {}),
+    (qml.BasisStatePreparation, (jnp.ones(3),), {"wires": [3, 2, 0]}),
+    (qml.BasisStatePreparation, (), {"basis_state": jnp.ones(1), "wires": [1]}),
+    (qml.CosineWindow, ([2, 3],), {}),
+    (qml.CosineWindow, (), {"wires": [2, 0, 1]}),
+    (qml.MottonenStatePreparation, (jnp.ones(4) / 2, [2, 3]), {}),
+    (
+        qml.MottonenStatePreparation,
+        (jnp.ones(8) / jnp.sqrt(8),),
+        {"wires": [3, 2, 0], "id": "your_id"},
+    ),
+    (qml.MottonenStatePreparation, (), {"state_vector": jnp.array([1.0, 0.0]), "wires": [1]}),
+    # Need to fix AllSinglesDoubles positional args: Currently have to pass hf_state as kwarg
+    # if we want to pass wires as kwarg, see issue #5521
+    (
+        qml.AllSinglesDoubles,
+        (jnp.ones(3), [2, 3, 0, 1]),
+        {
+            "singles": [[0, 2], [1, 3]],
+            "doubles": [[2, 3, 0, 1]],
+            "hf_state": np.array([0, 1, 1, 0]),
+        },
+    ),
+    (
+        qml.AllSinglesDoubles,
+        (jnp.zeros(3),),
+        {
+            "singles": [[0, 2], [1, 3]],
+            "doubles": [[2, 3, 0, 1]],
+            "wires": list(range(8)),
+            "hf_state": np.ones(8, dtype=int),
+        },
+    ),
+    # (qml.AllSinglesDoubles, (jnp.ones(3), [2, 3, 0, 1], np.array([0, 1, 1, 0])), {"singles": [[0, 2], [1, 3]], "doubles": [[2,3,0,1]]}), # Can't trace
+    (qml.AQFT, (1, [0, 1, 2]), {}),
+    (qml.AQFT, (2,), {"wires": [0, 1, 2, 3]}),
+    (qml.AQFT, (), {"order": 2, "wires": [0, 2, 3, 1]}),
+    (qml.QFT, ([0, 1],), {}),
+    (qml.QFT, (), {"wires": [0, 1]}),
+    (qml.ArbitraryUnitary, (jnp.ones(15), [2, 3]), {}),
+    (qml.ArbitraryUnitary, (jnp.zeros(15),), {"wires": [3, 2]}),
+    (qml.ArbitraryUnitary, (), {"weights": jnp.ones(3), "wires": [1]}),
+    (qml.FABLE, (jnp.eye(4), [2, 3, 0, 1, 5]), {}),
+    (qml.FABLE, (jnp.ones((4, 4)),), {"wires": [0, 3, 2, 1, 9]}),
+    (
+        qml.FABLE,
+        (),
+        {"input_matrix": jnp.array([[1, 1], [1, -1]]) / np.sqrt(2), "wires": [1, 10, 17]},
+    ),
+    (qml.FermionicSingleExcitation, (0.421,), {"wires": [0, 3, 2]}),
+    (qml.FlipSign, (7,), {"wires": [0, 3, 2]}),
+    (qml.FlipSign, (np.array([1, 0, 0]), [0, 1, 2]), {}),
+    (
+        qml.kUpCCGSD,
+        (jnp.ones((1, 6)), [0, 1, 2, 3]),
+        {"k": 1, "delta_sz": 0, "init_state": [1, 1, 0, 0]},
+    ),
+    (qml.Permute, (np.array([1, 2, 0]), [0, 1, 2]), {}),
+    (qml.Permute, (np.array([1, 2, 0]),), {"wires": [0, 1, 2]}),
+    (
+        qml.TwoLocalSwapNetwork,
+        ([0, 1, 2, 3, 4],),
+        {"acquaintances": lambda index, wires, param=None: qml.CNOT(index)},
+    ),
+    (qml.GroverOperator, (), {"wires": [0, 1]}),
+    (qml.GroverOperator, ([0, 1],), {}),
+    (
+        qml.UCCSD,
+        (jnp.ones(3), [2, 3, 0, 1]),
+        {"s_wires": [[0], [1]], "d_wires": [[[2], [3]]], "init_state": [0, 1, 1, 0]},
+    ),
+]
+
+
+@pytest.mark.parametrize("template, args, kwargs", unmodified_templates_cases)
+def test_unmodified_templates(template, args, kwargs):
+    """Test that templates with unmodified primitive binds are captured as expected."""
+
+    # Make sure the input data is valid
+    template(*args, **kwargs)
+
+    # Make sure the template actually is not modified in its primitive binding function
+    assert template._primitive_bind_call.__code__ == original_op_bind_code
+
+    def fn(*args):
+        template(*args, **kwargs)
+
+    jaxpr = jax.make_jaxpr(fn)(*args)
+
+    # Check basic structure of jaxpr: single equation with template primitive
+    assert len(jaxpr.eqns) == 1
+    eqn = jaxpr.eqns[0]
+    assert eqn.primitive == template._primitive
+
+    # Check that all arguments are passed correctly, taking wires parsing into account
+    # Also, store wires for later
+    if "wires" in kwargs:
+        # If wires are in kwargs, they are not invars to the jaxpr
+        num_invars_wo_wires = len(eqn.invars) - len(kwargs["wires"])
+        assert eqn.invars[:num_invars_wo_wires] == jaxpr.jaxpr.invars
+        wires = kwargs.pop("wires")
+    else:
+        # If wires are in args, they are also invars to the jaxpr
+        assert eqn.invars == jaxpr.jaxpr.invars
+        wires = args[-1]
+
+    # Check outvars; there should only be the DropVar returned by the template
+    assert len(eqn.outvars) == 1
+    assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+    # Check that `n_wires` is inferred correctly
+    if isinstance(wires, int):
+        wires = (wires,)
+    assert eqn.params.pop("n_wires") == len(wires)
+    # Check that remaining kwargs are passed properly to the eqn
+    assert eqn.params == kwargs
+
+
+# Only add a template to the following list if you manually added a test for it to
+# TestModifiedTemplates below.
+tested_modified_templates = [
+    qml.TrotterProduct,
+    qml.AmplitudeAmplification,
+    qml.ApproxTimeEvolution,
+    qml.BasisRotation,
+    qml.CommutingEvolution,
+    qml.ControlledSequence,
+    qml.FermionicDoubleExcitation,
+    qml.HilbertSchmidt,
+    qml.LocalHilbertSchmidt,
+    qml.QDrift,
+    qml.QSVT,
+    qml.QuantumMonteCarlo,
+    qml.QuantumPhaseEstimation,
+    qml.Qubitization,
+    qml.Reflection,
+    qml.Select,
+    qml.MERA,
+    qml.MPS,
+    qml.TTN,
+]
+
+
+class TestModifiedTemplates:
+    """Test that templates with custom primitive binds are captured as expected."""
+
+    @pytest.mark.parametrize(
+        "template, kwargs",
+        [
+            (qml.TrotterProduct, {"order": 2}),
+            (qml.ApproxTimeEvolution, {"n": 2}),
+            (qml.CommutingEvolution, {"frequencies": (1.2, 2)}),
+            (qml.QDrift, {"n": 2, "seed": 10}),
+        ],
+    )
+    def test_evolution_ops(self, template, kwargs):
+        """Test the primitive bind call of Hamiltonian time evolution templates."""
+
+        coeffs = [0.25, 0.75]
+        ops = [qml.X(0), qml.Z(0)]
+        H = qml.dot(coeffs, ops)
+
+        def qfunc(Hi):
+            template(Hi, time=2.4, **kwargs)
+
+        # Validate inputs
+        qfunc(H)
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)(H)
+
+        assert len(jaxpr.eqns) == 6
+
+        # due to flattening and unflattening H
+        assert jaxpr.eqns[0].primitive == qml.X._primitive
+        assert jaxpr.eqns[1].primitive == qml.ops.SProd._primitive
+        assert jaxpr.eqns[2].primitive == qml.Z._primitive
+        assert jaxpr.eqns[3].primitive == qml.ops.SProd._primitive
+        assert jaxpr.eqns[4].primitive == qml.ops.Sum._primitive
+
+        eqn = jaxpr.eqns[5]
+        assert eqn.primitive == template._primitive
+        assert eqn.invars == jaxpr.eqns[4].outvars  # the sum op
+
+        assert eqn.params == kwargs | {"time": 2.4}
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, coeffs[0], coeffs[1])
+
+        assert len(q) == 1
+        assert q.queue[0] == template(H, time=2.4, **kwargs)
+
+    def test_amplitude_amplification(self):
+        """Test the primitive bind call of AmplitudeAmplification."""
+
+        U = qml.Hadamard(0)
+        O = qml.FlipSign(1, 0)
+        iters = 3
+
+        kwargs = {"iters": iters, "fixed_point": False, "p_min": 0.4}
+
+        def qfunc(U, O):
+            qml.AmplitudeAmplification(U, O, **kwargs)
+
+        # Validate inputs
+        qfunc(U, O)
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)(U, O)
+
+        assert len(jaxpr.eqns) == 3
+
+        assert jaxpr.eqns[0].primitive == qml.Hadamard._primitive
+        assert jaxpr.eqns[1].primitive == qml.FlipSign._primitive
+
+        eqn = jaxpr.eqns[2]
+        assert eqn.primitive == qml.AmplitudeAmplification._primitive
+        assert eqn.invars[0] == jaxpr.eqns[0].outvars[0]  # Hadamard
+        assert eqn.invars[1] == jaxpr.eqns[1].outvars[0]  # FlipSign
+        assert eqn.params == kwargs
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)
+
+        assert len(q) == 1
+        assert q.queue[0] == qml.AmplitudeAmplification(U, O, **kwargs)
+
+    def test_basis_rotation(self):
+        """Test the primitive bind call of BasisRotation."""
+
+        mat = np.eye(4)
+        wires = [0, 5]
+
+        def qfunc(wires, mat):
+            qml.BasisRotation(wires, mat, check=True)
+
+        # Validate inputs
+        qfunc(wires, mat)
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)(wires, mat)
+
+        assert len(jaxpr.eqns) == 1
+
+        eqn = jaxpr.eqns[0]
+        assert eqn.primitive == qml.BasisRotation._primitive
+        assert eqn.invars == jaxpr.jaxpr.invars
+        assert eqn.params == {"check": True, "id": None}
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, *wires, mat)
+
+        assert len(q) == 1
+        assert q.queue[0] == qml.BasisRotation(wires=wires, unitary_matrix=mat, check=True)
+
+    def test_controlled_sequence(self):
+        """Test the primitive bind call of ControlledSequence."""
+
+        assert (
+            qml.ControlledSequence._primitive_bind_call.__code__
+            == qml.ops.op_math.SymbolicOp._primitive_bind_call.__code__
+        )
+
+        base = qml.RX(0.5, 0)
+        control = [1, 5]
+
+        def fn(base):
+            qml.ControlledSequence(base, control=control)
+
+        # Validate inputs
+        fn(base)
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(fn)(base)
+
+        assert len(jaxpr.eqns) == 2
+        assert jaxpr.eqns[0].primitive == qml.RX._primitive
+
+        eqn = jaxpr.eqns[1]
+        assert eqn.primitive == qml.ControlledSequence._primitive
+        assert eqn.invars == jaxpr.eqns[0].outvars
+        assert eqn.params == {"control": control}
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 0.5)
+
+        assert len(q) == 1  # One for each control
+        assert q.queue[0] == qml.ControlledSequence(base, control)
+
+    def test_fermionic_double_excitation(self):
+        """Test the primitive bind call of FermionicDoubleExcitation."""
+
+        weight = 0.251
+
+        kwargs = {"wires1": [0, 6], "wires2": [2, 3]}
+
+        def qfunc(weight):
+            qml.FermionicDoubleExcitation(weight, **kwargs)
+
+        # Validate inputs
+        qfunc(weight)
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)(weight)
+
+        assert len(jaxpr.eqns) == 1
+
+        eqn = jaxpr.eqns[0]
+        assert eqn.primitive == qml.FermionicDoubleExcitation._primitive
+        assert eqn.invars == jaxpr.jaxpr.invars
+        assert eqn.params == kwargs
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, weight)
+
+        assert len(q) == 1
+        assert q.queue[0] == qml.FermionicDoubleExcitation(weight, **kwargs)
+
+    @pytest.mark.parametrize("template", [qml.HilbertSchmidt, qml.LocalHilbertSchmidt])
+    def test_hilbert_schmidt(self, template):
+        """Test the primitive bind call of HilbertSchmidt and LocalHilbertSchmidt."""
+
+        v_params = np.array([0.6])
+
+        kwargs = {
+            "u_tape": qml.tape.QuantumScript([qml.Hadamard(0)]),
+            "v_function": lambda params: qml.RZ(params[0], wires=1),
+            "v_wires": [1],
+            "id": None,
+        }
+
+        def qfunc(v_params):
+            template(v_params, **kwargs)
+
+        # Validate inputs
+        qfunc(v_params)
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)(v_params)
+
+        assert len(jaxpr.eqns) == 1
+
+        eqn = jaxpr.eqns[0]
+        assert eqn.primitive == template._primitive
+        assert eqn.invars == jaxpr.jaxpr.invars
+        assert eqn.params == kwargs
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, v_params)
+
+        assert len(q) == 1
+        assert qml.equal(q.queue[0], template(v_params, **kwargs))
+
+    @pytest.mark.parametrize("template", [qml.MERA, qml.MPS, qml.TTN])
+    def test_tensor_networks(self, template):
+        """Test the primitive bind call of MERA, MPS, and TTN."""
+
+        def block(weights, wires):
+            return [
+                qml.CNOT(wires),
+                qml.RY(weights[0], wires[0]),
+                qml.RY(weights[1], wires[1]),
+            ]
+
+        wires = list(range(4))
+        n_block_wires = 2
+        n_blocks = template.get_n_blocks(wires, n_block_wires)
+
+        kwargs = {
+            "wires": wires,
+            "n_block_wires": n_block_wires,
+            "block": block,
+            "n_params_block": 2,
+            "template_weights": [[0.1, -0.3]] * n_blocks,
+        }
+
+        def qfunc():
+            template(**kwargs)
+
+        # Validate inputs
+        qfunc()
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)()
+
+        assert len(jaxpr.eqns) == 1
+
+        eqn = jaxpr.eqns[0]
+        assert eqn.primitive == template._primitive
+
+        expected_params = {
+            "n_block_wires": n_block_wires,
+            "block": block,
+            "n_params_block": 2,
+            "template_weights": kwargs["template_weights"],
+            "id": None,
+            "n_wires": 4,
+        }
+        if template is qml.MPS:
+            expected_params["offset"] = None
+        assert eqn.params == expected_params
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)
+
+        assert len(q) == 1
+        assert qml.equal(q.queue[0], template(**kwargs))
+
+    def test_qsvt(self):
+        """Test the primitive bind call of QSVT."""
+
+        A = np.array([[0.1]])
+        block_encode = qml.BlockEncode(A, wires=[0, 1])
+        shifts = [qml.PCPhase(i + 0.1, dim=1, wires=[0, 1]) for i in range(3)]
+
+        def qfunc(block_encode):
+            qml.QSVT(block_encode, projectors=shifts)
+
+        # Validate inputs
+        qfunc(block_encode)
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)(block_encode)
+
+        assert len(jaxpr.eqns) == 2
+
+        # due to flattening and unflattening BlockEncode
+        assert jaxpr.eqns[0].primitive == qml.BlockEncode._primitive
+
+        eqn = jaxpr.eqns[1]
+        assert eqn.primitive == qml.QSVT._primitive
+        assert eqn.invars == jaxpr.eqns[0].outvars
+        assert eqn.params == {"projectors": shifts}
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, A)
+
+        assert len(q) == 1
+        assert q.queue[0] == qml.QSVT(block_encode, shifts)
+
+    def test_quantum_monte_carlo(self):
+        """Test the primitive bind call of QuantumMonteCarlo."""
+
+        # This test follows the docstring example
+
+        from scipy.stats import norm
+
+        m = 5
+        M = 2**m
+        n = 10
+
+        xs = np.linspace(-np.pi, np.pi, M)
+        probs = np.array([norm().pdf(x) for x in xs])
+        probs /= np.sum(probs)
+
+        def func(i):
+            return np.sin(xs[i]) ** 2
+
+        target_wires = range(m + 1)
+        estimation_wires = range(m + 1, n + m + 1)
+
+        kwargs = {"func": func, "target_wires": target_wires, "estimation_wires": estimation_wires}
+
+        def qfunc(probs):
+            qml.QuantumMonteCarlo(probs, **kwargs)
+
+        # Validate inputs
+        qfunc(probs)
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)(probs)
+
+        assert len(jaxpr.eqns) == 1
+
+        eqn = jaxpr.eqns[0]
+        assert eqn.primitive == qml.QuantumMonteCarlo._primitive
+        assert eqn.invars == jaxpr.jaxpr.invars
+        assert eqn.params == kwargs
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, probs)
+
+        assert len(q) == 1
+        assert q.queue[0] == qml.QuantumMonteCarlo(probs, **kwargs)
+
+    def test_qubitization(self):
+        """Test the primitive bind call of Qubitization."""
+
+        hamiltonian = qml.dot([0.5, 1.2, -0.84], [qml.X(2), qml.Hadamard(3), qml.Z(2) @ qml.Y(3)])
+        kwargs = {"hamiltonian": hamiltonian, "control": [0, 1]}
+
+        def qfunc():
+            qml.Qubitization(**kwargs)
+
+        # Validate inputs
+        qfunc()
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)()
+
+        assert len(jaxpr.eqns) == 1
+
+        eqn = jaxpr.eqns[0]
+        assert eqn.primitive == qml.Qubitization._primitive
+        assert eqn.invars == jaxpr.jaxpr.invars
+        assert eqn.params == kwargs
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)
+
+        assert len(q) == 1
+        assert qml.equal(q.queue[0], qml.Qubitization(**kwargs))
+
+    @pytest.mark.parametrize(
+        "template, kwargs",
+        [
+            (qml.QuantumPhaseEstimation, {"estimation_wires": range(2, 4)}),
+            (qml.Reflection, {"alpha": np.pi / 2, "reflection_wires": [0]}),
+        ],
+    )
+    def test_quantum_phase_estimation_and_reflection(self, template, kwargs):
+        """Test the primitive bind call of QuantumPhaseEstimation and Reflection."""
+
+        op = qml.RX(np.pi / 2, 0) @ qml.Hadamard(1)
+
+        def qfunc(op):
+            template(op, **kwargs)
+
+        # Validate inputs
+        qfunc(op)
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)(op)
+
+        assert len(jaxpr.eqns) == 4
+
+        assert jaxpr.eqns[0].primitive == qml.RX._primitive
+        assert jaxpr.eqns[1].primitive == qml.Hadamard._primitive
+        assert jaxpr.eqns[2].primitive == qml.ops.op_math.Prod._primitive
+
+        eqn = jaxpr.eqns[3]
+        assert eqn.primitive == template._primitive
+        assert eqn.invars == jaxpr.eqns[2].outvars
+        assert eqn.params == kwargs
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, np.pi / 2)
+
+        assert len(q) == 1
+        assert qml.equal(q.queue[0], template(op, **kwargs))
+
+    def test_select(self):
+        """Test the primitive bind call of Select."""
+
+        ops = [qml.X(2), qml.RX(0.2, 3), qml.Y(2), qml.Z(3)]
+        kwargs = {"ops": ops, "control": [0, 1]}
+
+        def qfunc():
+            qml.Select(**kwargs)
+
+        # Validate inputs
+        qfunc()
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)()
+
+        assert len(jaxpr.eqns) == 1
+
+        eqn = jaxpr.eqns[0]
+        assert eqn.primitive == qml.Select._primitive
+        assert eqn.invars == jaxpr.jaxpr.invars
+        assert eqn.params == kwargs
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)
+
+        assert len(q) == 1
+        assert qml.equal(q.queue[0], qml.Select(**kwargs))
+
+
+def filter_fn(member: Any) -> bool:
+    """Determine whether a member of a module is a class and genuinely belongs to
+    qml.templates."""
+    return inspect.isclass(member) and member.__module__.startswith("pennylane.templates")
+
+
+_, all_templates = zip(*inspect.getmembers(qml.templates, filter_fn))
+
+unmodified_templates = [template for template, *_ in unmodified_templates_cases]
+unsupported_templates = [
+    qml.CVNeuralNetLayers,
+    qml.DisplacementEmbedding,
+    qml.Interferometer,
+    qml.QutritBasisStatePreparation,
+    qml.SqueezingEmbedding,
+]
+modified_templates = [
+    t for t in all_templates if t not in unmodified_templates + unsupported_templates
+]
+
+
+@pytest.mark.parametrize("template", modified_templates)
+def test_templates_are_modified(template):
+    """Test that all templates that are not listed as unmodified in the test cases above
+    actually have their _primitive_bind_call modified."""
+    # Make sure the template actually is modified in its primitive binding function
+    assert template._primitive_bind_call.__code__ != original_op_bind_code
+
+
+def test_all_modified_templates_are_tested():
+    """Test that all templates in `modified_templates` (automatically computed and
+    validated above) also are in `tested_modified_templates` (manually created and
+    expected to resemble all tested templates."""
+    assert set(modified_templates) == set(tested_modified_templates)
diff --git a/tests/templates/test_layers/test_gate_fabric.py b/tests/templates/test_layers/test_gate_fabric.py
index 02fd264ac..ae14dc173 100644
--- a/tests/templates/test_layers/test_gate_fabric.py
+++ b/tests/templates/test_layers/test_gate_fabric.py
@@ -753,7 +753,7 @@ class TestAttributes:
         """Test that the shape function warns if there are not enough qubits."""
 
         with pytest.raises(
-            ValueError, match="This template requires the number of qubits to be greater than four"
+            ValueError, match="This template requires the number of qubits to be at least four"
         ):
             qml.GateFabric.shape(3, 1)
 
diff --git a/tests/templates/test_subroutines/test_qpe.py b/tests/templates/test_subroutines/test_qpe.py
index 2139d263c..46b7c0d6b 100644
--- a/tests/templates/test_subroutines/test_qpe.py
+++ b/tests/templates/test_subroutines/test_qpe.py
@@ -24,6 +24,7 @@ import pennylane as qml
 def test_standard_validity():
     """Test standard validity criteria using assert_valid."""
     op = qml.QuantumPhaseEstimation(np.eye(4), target_wires=(0, 1), estimation_wires=[2, 5])
+    assert op.target_wires == qml.wires.Wires([0, 1])
     qml.ops.functions.assert_valid(op)
 
 
diff --git a/tests/templates/test_subroutines/test_select.py b/tests/templates/test_subroutines/test_select.py
index ab45a0b8a..808a31ce6 100644
--- a/tests/templates/test_subroutines/test_select.py
+++ b/tests/templates/test_subroutines/test_select.py
@@ -30,6 +30,7 @@ def test_standard_checks():
     control = [1]
 
     op = qml.Select(ops, control)
+    assert op.target_wires == qml.wires.Wires(0)
     qml.ops.functions.assert_valid(op)