Repository: PennyLaneAI/pennylane
Pull request #7750 (2025-06-25 16:27:47)

Move internal code of TransformProgram onto CotransformCache
**Context:**

Currently `TransformProgram` has a lot of code for calculating the classical jacobian and the argnums.  All this code is rather tied to the `CotransformCache` that holds onto the qnode, args, and kwargs.  

By moving this code onto the `CotransformCache`, it simplifies `TransformProgram` and opens up opportunities cleaning up `TransformProgram.__call__`.  But even if we don't clean up `TransformProgram.__call__`, this still simplifies `TransformProgram`.

**Description of the Change:**

- Extracts all the code that interacts with `CotransformCache` *into* `CotransformCache`.
- Cleans up imports in `transforms/core`

**Benefits:**

`TransformProgram` is more focused on it's real responsibilities.

**Possible Drawbacks:**

**Related GitHub Issues:**



Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 37f1f8476..4a05b13c4 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -859,6 +859,9 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 
 <h3>Internal changes ⚙️</h3>
 
+* Move private code in the `TransformProgram` onto the `CotransformCache` class.
+  [(#7750)](https://github.com/PennyLaneAI/pennylane/pull/7750)
+
 * Improve type hinting in the `workflow` module.
   [(#7745)](https://github.com/PennyLaneAI/pennylane/pull/7745)
 
diff --git a/pennylane/exceptions.py b/pennylane/exceptions.py
index 2bf9c97e8..62de34a26 100644
--- a/pennylane/exceptions.py
+++ b/pennylane/exceptions.py
@@ -34,3 +34,7 @@ class PennyLaneDeprecationWarning(UserWarning):  # pragma: no cover
 
 class ExperimentalWarning(UserWarning):  # pragma: no cover
     """Warning raised to indicate experimental/non-stable feature or support."""
+
+
+class TransformError(Exception):  # pragma: no cover
+    """Raised when there is an error with the transform logic."""
diff --git a/pennylane/ops/functions/eigvals.py b/pennylane/ops/functions/eigvals.py
index f7f5fd9c6..9c2889f68 100644
--- a/pennylane/ops/functions/eigvals.py
+++ b/pennylane/ops/functions/eigvals.py
@@ -21,8 +21,8 @@ import scipy
 
 import pennylane as qml
 from pennylane import transform
+from pennylane.exceptions import TransformError
 from pennylane.tape import QuantumScript, QuantumScriptBatch
-from pennylane.transforms import TransformError
 from pennylane.typing import PostprocessingFn, TensorLike
 
 
diff --git a/pennylane/ops/functions/matrix.py b/pennylane/ops/functions/matrix.py
index 78b7a0e7d..f73c5aab1 100644
--- a/pennylane/ops/functions/matrix.py
+++ b/pennylane/ops/functions/matrix.py
@@ -21,10 +21,10 @@ from typing import Union
 
 import pennylane as qml
 from pennylane import transform
+from pennylane.exceptions import TransformError
 from pennylane.operation import MatrixUndefinedError, Operator
 from pennylane.pauli import PauliSentence, PauliWord
 from pennylane.tape import QuantumScript, QuantumScriptBatch
-from pennylane.transforms import TransformError
 from pennylane.typing import PostprocessingFn, TensorLike
 
 
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index 00785ba50..5e9756061 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -601,7 +601,7 @@ class QuantumScript:
         return self._trainable_params
 
     @trainable_params.setter
-    def trainable_params(self, param_indices: list[int]):
+    def trainable_params(self, param_indices: list[int] | set[int]):
         """Store the indices of parameters that support differentiability.
 
         Args:
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index e958d7610..1470c732e 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -288,9 +288,10 @@ for a comprehensive overview of transforms and core functionalities, consult the
 
 # Leave as alias for backwards-compatibility
 from pennylane.tape import make_qscript as make_tape
+from pennylane.exceptions import TransformError
 
 # Import the decorators first to prevent circular imports when used in other transforms
-from .core import transform, TransformError
+from .core import transform
 from .batch_params import batch_params
 from .batch_input import batch_input
 from .batch_partial import batch_partial
diff --git a/pennylane/transforms/core/cotransform_cache.py b/pennylane/transforms/core/cotransform_cache.py
new file mode 100644
index 000000000..86c438d4d
--- /dev/null
+++ b/pennylane/transforms/core/cotransform_cache.py
@@ -0,0 +1,251 @@
+# Copyright 2023 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.
+"""
+This submodule contains the CotransformCache for handling the classical cotransform part of a transform.
+"""
+from functools import partial
+from typing import Optional
+
+import pennylane as qml  # for qml.workflow.construct_tape
+from pennylane import math
+from pennylane._grad import jacobian as autograd_jacobian
+from pennylane.exceptions import QuantumFunctionError
+from pennylane.typing import TensorLike
+
+from .transform_dispatcher import TransformContainer
+
+
+def _numpy_jac(*_, **__) -> TensorLike:
+    raise QuantumFunctionError("No trainable parameters.")
+
+
+def _autograd_jac(classical_function, argnums, *args, **kwargs) -> TensorLike:
+    if not math.get_trainable_indices(args) and argnums is None:
+        raise QuantumFunctionError("No trainable parameters.")
+    return autograd_jacobian(classical_function, argnum=argnums)(*args, **kwargs)
+
+
+# pylint: disable=import-outside-toplevel, unused-argument
+def _tf_jac(classical_function, argnums, *args, **kwargs) -> TensorLike:
+    if not math.get_trainable_indices(args):
+        raise QuantumFunctionError("No trainable parameters.")
+    import tensorflow as tf
+
+    with tf.GradientTape() as tape:
+        gate_params = classical_function(*args, **kwargs)
+    return tape.jacobian(gate_params, args)
+
+
+# pylint: disable=import-outside-toplevel, unused-argument
+def _torch_jac(classical_function, argnums, *args, **kwargs) -> TensorLike:
+    if not math.get_trainable_indices(args):
+        raise QuantumFunctionError("No trainable parameters.")
+    from torch.autograd.functional import jacobian
+
+    return jacobian(partial(classical_function, **kwargs), args)
+
+
+# pylint: disable=import-outside-toplevel
+def _jax_jac(classical_function, argnums, *args, **kwargs) -> TensorLike:
+    import jax
+
+    if argnums is None:
+        argnums = 0
+    return jax.jacobian(classical_function, argnums=argnums)(*args, **kwargs)
+
+
+_jac_map = {
+    None: _numpy_jac,
+    "numpy": _numpy_jac,
+    "autograd": _autograd_jac,
+    "tf": _tf_jac,
+    "torch": _torch_jac,
+    "jax": _jax_jac,
+    "jax-jit": _jax_jac,
+}
+
+
+# pylint: disable=unused-argument
+def _classical_preprocessing(qnode, program, tape_idx: int, *args, argnums=None, **kwargs):
+    """Returns the trainable gate parameters for a given QNode input.
+
+    While differentiating this again for each tape in the batch may be less efficient than desireable for large batches,
+    it cleanly works with all interfaces.
+    """
+    tape = qml.workflow.construct_tape(qnode, level=0)(*args, **kwargs)
+    tapes, _ = program((tape,))
+    return math.stack(tapes[tape_idx].get_parameters(trainable_only=True))
+
+
+def _jax_argnums_to_tape_trainable(qnode, argnums, program, args, kwargs):
+    """This function gets the tape parameters from the QNode construction given some argnums (only for Jax).
+    The tape parameters are transformed to JVPTracer if they are from argnums. This function imitates the behaviour
+    of Jax in order to mark trainable parameters.
+
+    Args:
+        qnode(qml.QNode): the quantum node.
+        argnums(int, list[int]): the parameters that we want to set as trainable (on the QNode level).
+        program(qml.transforms.core.TransformProgram): the transform program to be applied on the tape.
+
+    Return:
+        list[float, jax.JVPTracer]: List of parameters where the trainable one are `JVPTracer`.
+    """
+    import jax  # pylint: disable=import-outside-toplevel
+
+    tag = jax.core.TraceTag()
+    with jax.core.take_current_trace() as parent_trace:
+        trace = jax.interpreters.ad.JVPTrace(parent_trace, tag)
+        args_jvp = [
+            (
+                jax.interpreters.ad.JVPTracer(trace, arg, jax.numpy.zeros(arg.shape))
+                if i in argnums
+                else arg
+            )
+            for i, arg in enumerate(args)
+        ]
+        with jax.core.set_current_trace(trace):
+            tape = qml.workflow.construct_tape(qnode, level=0)(*args_jvp, **kwargs)
+            tapes, _ = program((tape,))
+
+    return tuple(tape.get_parameters(trainable_only=False) for tape in tapes)
+
+
+def _get_interface(qnode, args, kwargs) -> str:
+    if qnode.interface == "auto":
+        interface = math.get_interface(*args, *list(kwargs.values()))
+        try:
+            interface = math.get_canonical_interface_name(interface).value
+        except ValueError:
+            interface = "numpy"
+    else:
+        interface = qnode.interface
+    return interface
+
+
+class CotransformCache:
+    """A cache of the qnode, args, and kwargs that will can
+    be used to calculate the argnums and classical jacobian for the application
+    of a transform.
+
+    This class is an a implementation component for `TransformProgram`.
+    """
+
+    def __init__(self, qnode, args, kwargs):
+        self.qnode = qnode
+        self.args = args
+        self.kwargs = kwargs
+
+    def __eq__(self, other):
+        if not isinstance(other, CotransformCache):
+            return False
+        return self.qnode == other.qnode and self.args == other.args and self.kwargs == other.kwargs
+
+    @property
+    def _program(self):
+        return self.qnode.transform_program
+
+    def _get_idx_for_transform(self, transform):
+        for i, t in enumerate(self._program):
+            if t == transform:
+                return i
+        raise ValueError(f"Could not find {transform} in qnode's transform program.")
+
+    def get_classical_jacobian(self, transform: TransformContainer, tape_idx: int):
+        """Calculate the classical jacobian for a given transform.
+
+        Note that this function assumes that the transform exists at most one in the transform program.
+        Given transforms with classical cotransforms tend to be final transforms, this is a safe bet.
+
+        .. code-block:: python
+
+            @qml.gradients.param_shift
+            @qml.transforms.split_non_commuting
+            @qml.qnode(qml.device('default.qubit'))
+            def c(x, y):
+                qml.RX(2*x, 0)
+                qml.RY(x*y, 0)
+                return qml.expval(qml.Z(0)), qml.expval(qml.X(0))
+
+
+            ps_container = c.transform_program[-1]
+            x, y = qml.numpy.array(0.5), qml.numpy.array(3.0)
+
+            cc = CotransformCache(c, (x, y), {})
+
+        >>> cc.get_classical_jacobian(ps_container, tape_idx = 0)
+        (array([2., 3.]), array([0. , 0.5]))
+
+        """
+        transform_index = self._get_idx_for_transform(transform)
+        if not transform.classical_cotransform:
+            return None
+        argnums = self._program[-1].kwargs.get("argnums", None)  # pylint: disable=no-member
+
+        interface = _get_interface(self.qnode, self.args, self.kwargs)
+
+        subprogram = self._program[:transform_index]
+        f = partial(_classical_preprocessing, self.qnode, subprogram, tape_idx)
+        classical_jacobian = _jac_map[interface](f, argnums, *self.args, **self.kwargs)
+        return classical_jacobian
+
+    def get_argnums(self, transform: TransformContainer) -> Optional[list[set[int]]]:
+        """Calculate the trainable params from the argnums in the transform.
+
+        .. code-block:: python
+
+            @qml.transforms.split_non_commuting
+            @qml.qnode(qml.device('default.qubit'))
+            def c(x, y):
+                qml.RX(x[0], 0)
+                qml.RX(y, 0)
+                qml.RY(x[1], 0)
+                return qml.expval(qml.Z(0)), qml.expval(qml.X(0))
+
+            c = qml.gradients.param_shift(c, argnums=[0])
+
+            ps_container = c.transform_program[-1]
+            x, y = jax.numpy.array([0.5, 0.7]), jax.numpy.array(3.0)
+
+            cc = CotransformCache(c, (x, y), {})
+
+        >>> cc.get_argnums(ps_container)
+        [{0, 2}, {0, 2}]
+
+        """
+        transform_index = self._get_idx_for_transform(transform)
+        interface = _get_interface(self.qnode, self.args, self.kwargs)
+        if interface not in ["jax", "jax-jit"]:
+            return None
+
+        if "argnum" in self._program[transform_index].kwargs:
+            raise QuantumFunctionError(
+                "argnum does not work with the Jax interface. You should use argnums instead."
+            )
+
+        transform = self._program[transform_index]
+        argnums = self._program[-1].kwargs.get("argnums", None)  # pylint: disable=no-member
+
+        if argnums is None and math.get_interface(self.args[0]) != "jax":
+            raise QuantumFunctionError("No trainable parameters.")
+
+        argnums = [0] if argnums is None else argnums
+        argnums = [argnums] if isinstance(argnums, int) else argnums
+        # pylint: disable=protected-access
+        if (transform._use_argnum or transform.classical_cotransform) and argnums:
+            subprogram = self._program[:transform_index]
+            params = _jax_argnums_to_tape_trainable(
+                self.qnode, argnums, subprogram, self.args, self.kwargs
+            )
+            return [math.get_trainable_indices(param) for param in params]
+        return None
diff --git a/pennylane/transforms/core/transform_dispatcher.py b/pennylane/transforms/core/transform_dispatcher.py
index 8720bba9e..dee76d8df 100644
--- a/pennylane/transforms/core/transform_dispatcher.py
+++ b/pennylane/transforms/core/transform_dispatcher.py
@@ -22,13 +22,12 @@ from collections.abc import Sequence
 from copy import copy
 
 import pennylane as qml
+from pennylane import capture, math
+from pennylane.exceptions import TransformError
+from pennylane.queuing import AnnotatedQueue, QueuingManager, apply
 from pennylane.typing import ResultBatch
 
 
-class TransformError(Exception):
-    """Raised when there is an error with the transform logic."""
-
-
 def _create_plxpr_fallback_transform(tape_transform):
     # pylint: disable=import-outside-toplevel
     try:
@@ -40,7 +39,7 @@ def _create_plxpr_fallback_transform(tape_transform):
 
         def wrapper(*inner_args):
             tape = qml.tape.plxpr_to_tape(jaxpr, consts, *inner_args)
-            with qml.capture.pause():
+            with capture.pause():
                 tapes, _ = tape_transform(tape, *targs, **tkwargs)
 
             if len(tapes) > 1:
@@ -61,7 +60,7 @@ def _create_plxpr_fallback_transform(tape_transform):
 
             return tuple(out)
 
-        abstracted_axes, abstract_shapes = qml.capture.determine_abstracted_axes(args)
+        abstracted_axes, abstract_shapes = capture.determine_abstracted_axes(args)
         return jax.make_jaxpr(wrapper, abstracted_axes=abstracted_axes)(*abstract_shapes, *args)
 
     return plxpr_fallback_transform
@@ -257,8 +256,8 @@ class TransformDispatcher:  # pylint: disable=too-many-instance-attributes
                 return processing_fn(transformed_tapes)
             return transformed_tapes, processing_fn
 
-        if isinstance(obj, qml.QNode):
-            if qml.capture.enabled():
+        if isinstance(obj, qml.workflow.QNode):
+            if capture.enabled():
                 new_qnode = self.default_qnode_transform(obj, targs, tkwargs)
                 return self._capture_callable_transform(new_qnode, targs, tkwargs)
             return self._qnode_transform(obj, targs, tkwargs)
@@ -274,7 +273,7 @@ class TransformDispatcher:  # pylint: disable=too-many-instance-attributes
             )
 
         if callable(obj):
-            if qml.capture.enabled():
+            if capture.enabled():
                 return self._capture_callable_transform(obj, targs, tkwargs)
             return self._qfunc_transform(obj, targs, tkwargs)
 
@@ -395,7 +394,7 @@ class TransformDispatcher:  # pylint: disable=too-many-instance-attributes
         def qfunc_transformed(*args, **kwargs):
             import jax  # pylint: disable=import-outside-toplevel
 
-            flat_qfunc = qml.capture.flatfn.FlatFn(qfunc)
+            flat_qfunc = capture.flatfn.FlatFn(qfunc)
             jaxpr = jax.make_jaxpr(functools.partial(flat_qfunc, **kwargs))(*args)
             flat_args = jax.tree_util.tree_leaves(args)
 
@@ -426,11 +425,11 @@ class TransformDispatcher:  # pylint: disable=too-many-instance-attributes
 
         @functools.wraps(qfunc)
         def qfunc_transformed(*args, **kwargs):
-            with qml.queuing.AnnotatedQueue() as q:
+            with AnnotatedQueue() as q:
                 qfunc_output = qfunc(*args, **kwargs)
 
             tape = qml.tape.QuantumScript.from_queue(q)
-            with qml.QueuingManager.stop_recording():
+            with QueuingManager.stop_recording():
                 transformed_tapes, processing_fn = self._transform(tape, *targs, **tkwargs)
 
             if len(transformed_tapes) != 1:
@@ -445,7 +444,7 @@ class TransformDispatcher:  # pylint: disable=too-many-instance-attributes
                 return processing_fn(transformed_tapes)
 
             for op in transformed_tape.circuit:
-                qml.apply(op)
+                apply(op)
 
             mps = transformed_tape.measurements
 
@@ -458,8 +457,8 @@ class TransformDispatcher:  # pylint: disable=too-many-instance-attributes
             if isinstance(qfunc_output, (tuple, list)):
                 return type(qfunc_output)(mps)
 
-            interface = qml.math.get_interface(qfunc_output)
-            return qml.math.asarray(mps, like=interface)
+            interface = math.get_interface(qfunc_output)
+            return math.asarray(mps, like=interface)
 
         return qfunc_transformed
 
@@ -634,7 +633,7 @@ def _create_transform_primitive(name):
     ):  # pylint: disable=unused-argument
         args = all_args[args_slice]
         consts = all_args[consts_slice]
-        return qml.capture.eval_jaxpr(inner_jaxpr, consts, *args)
+        return capture.eval_jaxpr(inner_jaxpr, consts, *args)
 
     @transform_prim.def_abstract_eval
     def _(*_, inner_jaxpr, **__):
diff --git a/pennylane/transforms/core/transform_program.py b/pennylane/transforms/core/transform_program.py
index 9d5720932..9d9f02352 100644
--- a/pennylane/transforms/core/transform_program.py
+++ b/pennylane/transforms/core/transform_program.py
@@ -14,131 +14,16 @@
 """
 This module contains the ``TransformProgram`` class.
 """
-from collections import namedtuple
 from collections.abc import Sequence
 from functools import partial
 from typing import Optional, Union, overload
 
-import pennylane as qml
-from pennylane import math
-from pennylane.exceptions import QuantumFunctionError
+from pennylane.exceptions import TransformError
 from pennylane.tape import QuantumScriptBatch
 from pennylane.typing import BatchPostprocessingFn, PostprocessingFn, ResultBatch
 
-from .transform_dispatcher import TransformContainer, TransformDispatcher, TransformError
-
-CotransformCache = namedtuple("CotransformCache", ("qnode", "args", "kwargs"))
-
-
-def _get_interface(qnode, args, kwargs) -> str:
-    if qnode.interface == "auto":
-        interface = math.get_interface(*args, *list(kwargs.values()))
-        try:
-            interface = math.get_canonical_interface_name(interface).value
-        except ValueError:
-            interface = "numpy"
-    else:
-        interface = qnode.interface
-    return interface
-
-
-def _numpy_jac(*_, **__) -> qml.typing.TensorLike:
-    raise QuantumFunctionError("No trainable parameters.")
-
-
-def _autograd_jac(classical_function, argnums, *args, **kwargs) -> qml.typing.TensorLike:
-    if not math.get_trainable_indices(args) and argnums is None:
-        raise QuantumFunctionError("No trainable parameters.")
-    return qml.jacobian(classical_function, argnum=argnums)(*args, **kwargs)
-
-
-# pylint: disable=import-outside-toplevel, unused-argument
-def _tf_jac(classical_function, argnums, *args, **kwargs) -> qml.typing.TensorLike:
-    if not math.get_trainable_indices(args):
-        raise QuantumFunctionError("No trainable parameters.")
-    import tensorflow as tf
-
-    with tf.GradientTape() as tape:
-        gate_params = classical_function(*args, **kwargs)
-    return tape.jacobian(gate_params, args)
-
-
-# pylint: disable=import-outside-toplevel, unused-argument
-def _torch_jac(classical_function, argnums, *args, **kwargs) -> qml.typing.TensorLike:
-    if not math.get_trainable_indices(args):
-        raise QuantumFunctionError("No trainable parameters.")
-    from torch.autograd.functional import jacobian
-
-    return jacobian(partial(classical_function, **kwargs), args)
-
-
-# pylint: disable=import-outside-toplevel
-def _jax_jac(classical_function, argnums, *args, **kwargs) -> qml.typing.TensorLike:
-    import jax
-
-    if argnums is None:
-        argnums = 0
-    return jax.jacobian(classical_function, argnums=argnums)(*args, **kwargs)
-
-
-_jac_map = {
-    None: _numpy_jac,
-    "numpy": _numpy_jac,
-    "autograd": _autograd_jac,
-    "tf": _tf_jac,
-    "torch": _torch_jac,
-    "jax": _jax_jac,
-    "jax-jit": _jax_jac,
-}
-
-
-# pylint: disable=unused-argument
-def _classical_preprocessing(qnode, program, tape_idx: int, *args, **kwargs):
-    """Returns the trainable gate parameters for a given QNode input.
-
-    While differentiating this again for each tape in the batch may be less efficient than desireable for large batches,
-    it cleanly works with all interfaces.
-    """
-    tape = qml.workflow.construct_tape(qnode, level=0)(*args, **kwargs)
-    tapes, _ = program((tape,))
-    new_tape = tapes[tape_idx]
-    # Recalculate trainable parameters because the program is not obligated to maintain them
-    params = new_tape.get_parameters(trainable_only=False)
-    new_tape.trainable_params = qml.math.get_trainable_indices(params)
-    return math.stack(new_tape.get_parameters(trainable_only=True))
-
-
-def _jax_argnums_to_tape_trainable(qnode, argnums, program, args, kwargs):
-    """This function gets the tape parameters from the QNode construction given some argnums (only for Jax).
-    The tape parameters are transformed to JVPTracer if they are from argnums. This function imitates the behaviour
-    of Jax in order to mark trainable parameters.
-
-    Args:
-        qnode(qml.QNode): the quantum node.
-        argnums(int, list[int]): the parameters that we want to set as trainable (on the QNode level).
-        program(qml.transforms.core.TransformProgram): the transform program to be applied on the tape.
-
-    Return:
-        list[float, jax.JVPTracer]: List of parameters where the trainable one are `JVPTracer`.
-    """
-    import jax  # pylint: disable=import-outside-toplevel
-
-    tag = jax.core.TraceTag()
-    with jax.core.take_current_trace() as parent_trace:
-        trace = jax.interpreters.ad.JVPTrace(parent_trace, tag)
-        args_jvp = [
-            (
-                jax.interpreters.ad.JVPTracer(trace, arg, jax.numpy.zeros(arg.shape))
-                if i in argnums
-                else arg
-            )
-            for i, arg in enumerate(args)
-        ]
-        with jax.core.set_current_trace(trace):
-            tape = qml.workflow.construct_tape(qnode, level=0)(*args_jvp, **kwargs)
-            tapes, _ = program((tape,))
-
-    return tuple(tape.get_parameters(trainable_only=False) for tape in tapes)
+from .cotransform_cache import CotransformCache
+from .transform_dispatcher import TransformContainer, TransformDispatcher
 
 
 def _batch_postprocessing(
@@ -516,48 +401,6 @@ class TransformProgram:
             i -= 1
         return found
 
-    def _get_classical_jacobian(self, index: int, tape_idx: int):
-        if self.cotransform_cache is None or not self[index].classical_cotransform:
-            return None
-        argnums = self[-1].kwargs.get("argnums", None)  # pylint: disable=no-member
-        qnode, args, kwargs = self.cotransform_cache
-
-        interface = _get_interface(qnode, args, kwargs)
-
-        f = partial(_classical_preprocessing, qnode, self[:index], tape_idx)
-        classical_jacobian = _jac_map[interface](f, argnums, *args, **kwargs)
-        return classical_jacobian
-
-    def _get_argnums(self, index):
-        """It can be used inside the QNode to set all argnums (tape level) using argnums from the argnums at the QNode
-        level.
-        """
-        if self.cotransform_cache is None:
-            return None
-
-        qnode, args, kwargs = self.cotransform_cache
-        interface = _get_interface(qnode, args, kwargs)
-        if interface not in ["jax", "jax-jit"]:
-            return None
-
-        if "argnum" in self[index].kwargs:
-            raise QuantumFunctionError(
-                "argnum does not work with the Jax interface. You should use argnums instead."
-            )
-
-        transform = self[index]
-        argnums = self[-1].kwargs.get("argnums", None)  # pylint: disable=no-member
-
-        if argnums is None and math.get_interface(args[0]) != "jax":
-            raise QuantumFunctionError("No trainable parameters.")
-
-        argnums = [0] if argnums is None else argnums
-        # pylint: disable=protected-access
-        if (transform._use_argnum or transform.classical_cotransform) and argnums:
-            params = _jax_argnums_to_tape_trainable(qnode, argnums, self[:index], args, kwargs)
-            return [math.get_trainable_indices(param) for param in params]
-        return None
-
     def __call_tapes(
         self, tapes: QuantumScriptBatch
     ) -> tuple[QuantumScriptBatch, BatchPostprocessingFn]:
@@ -566,7 +409,7 @@ class TransformProgram:
 
         processing_fns_stack = []
 
-        for transform_idx, transform_container in enumerate(self):
+        for transform_container in self:
             transform, targs, tkwargs, cotransform, _, _, _ = transform_container
             tkwargs = {
                 key: value for key, value in tkwargs.items() if key not in {"argnums", "hybrid"}
@@ -574,11 +417,16 @@ class TransformProgram:
             execution_tapes, fns, slices, classical_fns = [], [], [], []
 
             start = 0
-            argnums = self._get_argnums(transform_idx)
+            argnums = (
+                self.cotransform_cache.get_argnums(transform_container)
+                if self.cotransform_cache
+                else None
+            )
 
             classical_jacobians = []
             for tape_idx, tape in enumerate(tapes):
                 if argnums is not None:
+                    # pylint: disable=unsubscriptable-object
                     tape.trainable_params = argnums[tape_idx]
                 new_tapes, fn = transform(tape, *targs, **tkwargs)
                 execution_tapes.extend(new_tapes)
@@ -588,7 +436,12 @@ class TransformProgram:
                 slices.append(slice(start, end))
                 start = end
 
-                classical_jacobians.append(self._get_classical_jacobian(transform_idx, tape_idx))
+                jac = (
+                    self.cotransform_cache.get_classical_jacobian(transform_container, tape_idx)
+                    if self.cotransform_cache
+                    else None
+                )
+                classical_jacobians.append(jac)
                 if cotransform and classical_jacobians[-1] is not None:
                     classical_fns.append(
                         partial(cotransform, cjac=classical_jacobians[-1], tape=tape)
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 3f1a4eb2e..bdaed5471 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -19,6 +19,7 @@ from typing import Callable, Optional, Sequence, Union
 from warnings import warn
 
 import pennylane as qml
+from pennylane.exceptions import TransformError
 from pennylane.measurements import (
     CountsMP,
     MeasurementValue,
@@ -30,7 +31,7 @@ from pennylane.measurements import (
 from pennylane.ops.op_math import ctrl
 from pennylane.queuing import QueuingManager
 from pennylane.tape import QuantumScript, QuantumScriptBatch
-from pennylane.transforms import TransformError, transform
+from pennylane.transforms import transform
 from pennylane.typing import PostprocessingFn
 from pennylane.wires import Wires
 
diff --git a/pennylane/transforms/optimization/merge_amplitude_embedding.py b/pennylane/transforms/optimization/merge_amplitude_embedding.py
index af129ba27..4aa6aa775 100644
--- a/pennylane/transforms/optimization/merge_amplitude_embedding.py
+++ b/pennylane/transforms/optimization/merge_amplitude_embedding.py
@@ -19,11 +19,11 @@ from typing import Sequence
 
 import pennylane as qml
 from pennylane import AmplitudeEmbedding
-from pennylane.exceptions import DeviceError
+from pennylane.exceptions import DeviceError, TransformError
 from pennylane.math import flatten, is_abstract, reshape
 from pennylane.queuing import QueuingManager
 from pennylane.tape import QuantumScript, QuantumScriptBatch
-from pennylane.transforms.core import TransformError, transform
+from pennylane.transforms.core import transform
 from pennylane.typing import PostprocessingFn
 
 
diff --git a/pennylane/transforms/zx/converter.py b/pennylane/transforms/zx/converter.py
index f56ee639f..688065962 100644
--- a/pennylane/transforms/zx/converter.py
+++ b/pennylane/transforms/zx/converter.py
@@ -21,10 +21,10 @@ from functools import partial
 import numpy as np
 
 import pennylane as qml
-from pennylane.exceptions import QuantumFunctionError
+from pennylane.exceptions import QuantumFunctionError, TransformError
 from pennylane.operation import Operator
 from pennylane.tape import QuantumScript, QuantumScriptBatch
-from pennylane.transforms import TransformError, transform
+from pennylane.transforms import transform
 from pennylane.typing import PostprocessingFn
 from pennylane.wires import Wires
 


Test patch:
diff --git a/tests/transforms/core/test_cotransform_cache.py b/tests/transforms/core/test_cotransform_cache.py
new file mode 100644
index 000000000..631c80b15
--- /dev/null
+++ b/tests/transforms/core/test_cotransform_cache.py
@@ -0,0 +1,198 @@
+# Copyright 2025 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.
+"""
+Tests for the CotransformCache object.
+"""
+import pytest
+
+import pennylane as qml
+from pennylane.transforms.core import TransformContainer
+from pennylane.transforms.core.cotransform_cache import CotransformCache
+
+
+def test_classical_jacobian_error_if_not_in_cache():
+    """Test a ValueError is raised if we request the classical jacobian for a transform not in the cotransform cache."""
+
+    @qml.qnode(qml.device("default.qubit"))
+    def c(x, y):
+        qml.RX(2 * x, 0)
+        qml.RY(x * y, 0)
+        return qml.expval(qml.Z(0)), qml.expval(qml.X(0))
+
+    container = TransformContainer(qml.transforms.merge_rotations.transform)
+    x, y = qml.numpy.array(0.5), qml.numpy.array(3.0)
+
+    cc = CotransformCache(c, (x, y), {})
+    with pytest.raises(ValueError, match=r"Could not find"):
+        cc.get_classical_jacobian(container, 0)
+
+
+def test_no_classical_jacobian_if_no_cotransform():
+    """Test that the classical jacobian is None if the transform does not have a cotransform."""
+
+    @qml.transforms.split_non_commuting
+    @qml.qnode(qml.device("default.qubit"))
+    def c(x, y):
+        qml.RX(2 * x, 0)
+        qml.RY(x * y, 0)
+        return qml.expval(qml.Z(0)), qml.expval(qml.X(0))
+
+    container = c.transform_program[-1]
+    x, y = qml.numpy.array(0.5), qml.numpy.array(3.0)
+
+    cc = CotransformCache(c, (x, y), {})
+    assert cc.get_classical_jacobian(container, 0) is None
+
+
+def test_simple_classical_jacobian():
+    """Test the calculation of a simple classical jacobian."""
+
+    @qml.gradients.param_shift
+    @qml.transforms.split_non_commuting
+    @qml.qnode(qml.device("default.qubit"))
+    def c(x, y):
+        qml.RX(2 * x, 0)
+        qml.RY(x * y, 0)
+        return qml.expval(qml.Z(0)), qml.expval(qml.X(0))
+
+    ps_container = c.transform_program[-1]
+    x, y = qml.numpy.array(0.5), qml.numpy.array(3.0)
+
+    a = CotransformCache(c, (x, y), {})
+    for i in range(2):
+        x_jac, y_jac = a.get_classical_jacobian(ps_container, i)
+        assert qml.math.allclose(x_jac, qml.numpy.array([2.0, 3.0]))
+        assert qml.math.allclose(y_jac, qml.numpy.array([0.0, 0.5]))
+
+
+@pytest.mark.jax
+@pytest.mark.parametrize(
+    "argnums, trainable_params",
+    [
+        (0, [{0, 2}, {0, 2}]),
+        (None, [{0, 2}, {0, 2}]),
+        ([0], [{0, 2}, {0, 2}]),
+        (
+            1,
+            [
+                {
+                    1,
+                },
+                {
+                    1,
+                },
+            ],
+        ),
+        (
+            [1],
+            [
+                {
+                    1,
+                },
+                {
+                    1,
+                },
+            ],
+        ),
+        ([0, 1], [{0, 1, 2}, {0, 1, 2}]),
+    ],
+)
+def test_simple_argnums(argnums, trainable_params):
+    """Test a simple calculation of the argnums."""
+
+    import jax
+
+    @qml.transforms.split_non_commuting
+    @qml.qnode(qml.device("default.qubit"))
+    def c(x, y):
+        qml.RX(x[0], 0)
+        qml.RX(y, 0)
+        qml.RY(x[1], 0)
+        return qml.expval(qml.Z(0)), qml.expval(qml.X(0))
+
+    c = qml.gradients.param_shift(c, argnums=argnums)
+
+    ps_container = c.transform_program[-1]
+    x, y = jax.numpy.array([0.5, 0.7]), jax.numpy.array(3.0)
+
+    cc = CotransformCache(c, (x, y), {})
+
+    assert cc.get_argnums(ps_container) == trainable_params
+
+
+@pytest.mark.jax
+def test_no_jax_argnum_error():
+    """Test that an error is raised in the jax interface is used with argnum"""
+
+    import jax
+
+    @qml.qnode(qml.device("default.qubit"))
+    def c(x, y):
+        qml.RX(x[0], 0)
+        qml.RX(y, 0)
+        qml.RY(x[1], 0)
+        return qml.expval(qml.Z(0))
+
+    c = qml.gradients.param_shift(c, argnum=[0])
+
+    ps_container = c.transform_program[-1]
+    x, y = jax.numpy.array([0.5, 0.7]), jax.numpy.array(3.0)
+
+    cc = CotransformCache(c, (x, y), {})
+
+    with pytest.raises(
+        qml.exceptions.QuantumFunctionError, match="argnum does not work with the Jax interface"
+    ):
+        cc.get_argnums(ps_container)
+
+
+@pytest.mark.jax
+def test_no_argnums_if_no_classical_cotransform():
+    """Test argnums is None if there is no classical cotransform."""
+    import jax
+
+    @qml.qnode(qml.device("default.qubit"))
+    def c(x, y):
+        qml.RX(x[0], 0)
+        qml.RX(y, 0)
+        qml.RY(x[1], 0)
+        return qml.expval(qml.Z(0))
+
+    c = qml.transforms.merge_rotations(c)
+
+    container = c.transform_program[-1]
+    x, y = jax.numpy.array([0.5, 0.7]), jax.numpy.array(3.0)
+
+    cc = CotransformCache(c, (x, y), {})
+
+    assert cc.get_argnums(container) is None
+
+
+def test_no_argnums_nonjax_interface():
+    """Test that the trainable params are None if the interface isn't jax."""
+
+    @qml.qnode(qml.device("default.qubit"))
+    def c(x, y):
+        qml.RX(x[0], 0)
+        qml.RX(y, 0)
+        qml.RY(x[1], 0)
+        return qml.expval(qml.Z(0))
+
+    c = qml.gradients.param_shift(c, argnum=[0])
+
+    ps_container = c.transform_program[-1]
+    x, y = qml.numpy.array([0.5, 0.7]), qml.numpy.array(3.0)
+
+    cc = CotransformCache(c, (x, y), {})
+    assert cc.get_argnums(ps_container) is None
diff --git a/tests/transforms/core/test_transform_program.py b/tests/transforms/core/test_transform_program.py
index 2e8671bb7..ae690a5b5 100644
--- a/tests/transforms/core/test_transform_program.py
+++ b/tests/transforms/core/test_transform_program.py
@@ -608,13 +608,12 @@ class TestClassicalCotransfroms:
             qml.RX(x, 0)
             return qml.expval(qml.Z(0))
 
-        program = TransformProgram()
-        program.add_transform(qml.gradients.param_shift, hybrid=True)
-        program.set_classical_component(circuit, (arg,), {})
+        circuit = qml.gradients.param_shift(circuit, hybrid=True)
+        circuit.transform_program.set_classical_component(circuit, (arg,), {})
 
         tape = qml.tape.QuantumScript([], [])
         with pytest.raises(QuantumFunctionError, match="No trainable parameters"):
-            program((tape,))
+            circuit.transform_program((tape,))
 
 
 class TestTransformProgramCall: