Repository: PennyLaneAI/pennylane
Pull request #7418 (2025-05-12 19:58:26)

[circ. deps] enforce `optimize` as `auxiliary` module
**Context:**

The goal of this PR is to isolate the `optimize` module as an `auxiliary` module.

**Description of the Change:**

- Updates `tach.toml` accordingly
- Fixes any errors `tach` identifies
- Uses `ops.LinearCombination` rather than `qml.Hamiltonian` in `riemannian_gradient.py`

⚠️ This module depends on `fourier` (a _tertiary_ module). This dependency is not easy to disentangle, so for now we are ignoring the import error and a follow up story sc-89244 has been created.

**Benefits:** Module dependency enforcement.

**Possible Drawbacks:** None identified.

[sc-88817]


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.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7418?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.69%. Comparing base [(`fadcab6`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/fadcab6115d591b1d26758d7761d26ce1a657bd1?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`fa41744`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/fa41744d03f8ba80b1ee3faee5fb2a8a86f6a508?dropdown=coverage&el=desc&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    #7418   +/-   ##
=======================================
  Coverage   99.69%   99.69%           
=======================================
  Files         529      529           
  Lines       50672    50694   +22     
=======================================
+ Hits        50517    50539   +22     
  Misses        155      155           
```

</details>

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

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
andrijapau: CC @dwierichs @JerryChen97 - reverted the forcing of kwargs in this PR as,
1. they are beyond the scope of this PR
2. will be addressed this quarter when we update the PyLint version to latest
3. need input from product.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 0a1521ed2..d16c79669 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -272,6 +272,9 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 
 <h3>Internal changes ⚙️</h3>
 
+* Enforce `optimize` module to be an auxiliary layer module.
+  [(#7418)](https://github.com/PennyLaneAI/pennylane/pull/7418)
+
 * A `RuntimeWarning` raised when using versions of JAX > 0.4.28 has been removed.
   [(#7398)](https://github.com/PennyLaneAI/pennylane/pull/7398)
 
diff --git a/pennylane/optimize/adaptive.py b/pennylane/optimize/adaptive.py
index de57b1b33..8edca18f2 100644
--- a/pennylane/optimize/adaptive.py
+++ b/pennylane/optimize/adaptive.py
@@ -15,11 +15,15 @@
 import copy
 
 # pylint: disable= no-value-for-parameter, protected-access, not-callable
-import pennylane as qml
+from pennylane import math
 from pennylane import numpy as pnp
-from pennylane import transform
+from pennylane._grad import grad
 from pennylane.tape import QuantumScript, QuantumScriptBatch
+from pennylane.transforms.core import transform
 from pennylane.typing import PostprocessingFn
+from pennylane.workflow import construct_tape
+
+from .gradient_descent import GradientDescentOptimizer
 
 
 @transform
@@ -199,7 +203,7 @@ class AdaptiveOptimizer:
         """
         cost = circuit()
         qnode = copy.copy(circuit)
-        tape = qml.workflow.construct_tape(qnode)()
+        tape = construct_tape(qnode)()
 
         if drain_pool:
             operator_pool = [
@@ -213,10 +217,10 @@ class AdaptiveOptimizer:
 
         params = pnp.array([gate.parameters[0] for gate in operator_pool], requires_grad=True)
         qnode.func = self._circuit
-        grads = qml.grad(qnode)(params, gates=operator_pool, initial_circuit=circuit.func)
+        grads = grad(qnode)(params, gates=operator_pool, initial_circuit=circuit.func)
 
         selected_gates = [operator_pool[pnp.argmax(abs(grads))]]
-        optimizer = qml.GradientDescentOptimizer(stepsize=self.stepsize)
+        optimizer = GradientDescentOptimizer(stepsize=self.stepsize)
 
         if params_zero:
             params = pnp.zeros(len(selected_gates))
@@ -230,4 +234,4 @@ class AdaptiveOptimizer:
 
         qnode.func = append_gate(circuit.func, params, selected_gates)
 
-        return qnode, cost, max(abs(qml.math.toarray(grads)))
+        return qnode, cost, max(abs(math.toarray(grads)))
diff --git a/pennylane/optimize/qng.py b/pennylane/optimize/qng.py
index 44778d9bd..e01463881 100644
--- a/pennylane/optimize/qng.py
+++ b/pennylane/optimize/qng.py
@@ -15,21 +15,23 @@
 import numbers
 from collections.abc import Iterable
 
-import pennylane as qml
-
 # pylint: disable=too-many-branches
 # pylint: disable=too-many-arguments
+from pennylane import math
 from pennylane import numpy as pnp
+from pennylane.gradients.metric_tensor import metric_tensor
+from pennylane.wires import Wires
+from pennylane.workflow import QNode
 
 from .gradient_descent import GradientDescentOptimizer
 
 
 def _reshape_and_regularize(tensor, lam):
-    shape = qml.math.shape(tensor)
-    size = 1 if shape == () else qml.math.prod(shape[: len(shape) // 2])
-    tensor = qml.math.reshape(tensor, (size, size))
+    shape = math.shape(tensor)
+    size = 1 if shape == () else math.prod(shape[: len(shape) // 2])
+    tensor = math.reshape(tensor, (size, size))
     # Add regularization
-    tensor += lam * qml.math.eye(size, like=tensor)
+    tensor += lam * math.eye(size, like=tensor)
     return tensor
 
 
@@ -198,7 +200,7 @@ class QNGOptimizer(GradientDescentOptimizer):
             prior to the step
         """
         # pylint: disable=arguments-differ
-        if not isinstance(qnode, qml.QNode) and metric_tensor_fn is None:
+        if not isinstance(qnode, QNode) and metric_tensor_fn is None:
             raise ValueError(
                 "The objective function must be encoded as a single QNode for the natural gradient "
                 "to be automatically computed. Otherwise, metric_tensor_fn must be explicitly "
@@ -207,7 +209,7 @@ class QNGOptimizer(GradientDescentOptimizer):
 
         if recompute_tensor or self.metric_tensor is None:
             if metric_tensor_fn is None:
-                metric_tensor_fn = qml.metric_tensor(qnode, approx=self.approx)
+                metric_tensor_fn = metric_tensor(qnode, approx=self.approx)
 
             mt = metric_tensor_fn(*args, **kwargs)
             if isinstance(mt, tuple):
@@ -304,7 +306,7 @@ def _flatten_np(x):
         yield from _flatten_np(
             x.flat
         )  # should we allow object arrays? or just "yield from x.flat"?
-    elif isinstance(x, qml.wires.Wires):
+    elif isinstance(x, Wires):
         # Reursive calls to flatten `Wires` will cause infinite recursion (`Wires` atoms are `Wires`).
         # Since Wires are always flat, just yield.
         yield from x
diff --git a/pennylane/optimize/qnspsa.py b/pennylane/optimize/qnspsa.py
index ad44618fa..0dbd33a97 100644
--- a/pennylane/optimize/qnspsa.py
+++ b/pennylane/optimize/qnspsa.py
@@ -16,8 +16,11 @@ import warnings
 
 from scipy.linalg import sqrtm
 
-import pennylane as qml
 from pennylane import numpy as pnp
+from pennylane.measurements import probs
+from pennylane.ops import adjoint
+from pennylane.tape import QuantumScript
+from pennylane.workflow import construct_tape, execute
 
 
 class QNSPSAOptimizer:
@@ -126,6 +129,7 @@ class QNSPSAOptimizer:
 
     # pylint: disable=too-many-arguments
     # pylint: disable=too-many-instance-attributes
+    # pylint: disable=too-many-positional-arguments
     def __init__(
         self,
         stepsize=1e-3,
@@ -222,7 +226,7 @@ class QNSPSAOptimizer:
             all_tensor_dirs.append(tensor_dirs)
 
         # nosemgrep
-        raw_results = qml.execute(all_grad_tapes + all_metric_tapes, cost.device)
+        raw_results = execute(all_grad_tapes + all_metric_tapes, cost.device)
         grads = [
             self._post_process_grad(raw_results[2 * i : 2 * i + 2], all_grad_dirs[i])
             for i in range(self.resamplings)
@@ -357,9 +361,9 @@ class QNSPSAOptimizer:
             args_plus[index] = arg + self.finite_diff_step * direction
             args_minus[index] = arg - self.finite_diff_step * direction
 
-        tape = qml.workflow.construct_tape(cost)(*args_plus, **kwargs)
+        tape = construct_tape(cost)(*args_plus, **kwargs)
         tape_plus = tape.copy(copy_operations=True)
-        tape = qml.workflow.construct_tape(cost)(*args_minus, **kwargs)
+        tape = construct_tape(cost)(*args_minus, **kwargs)
         tape_minus = tape.copy(copy_operations=True)
         return [tape_plus, tape_minus], dirs
 
@@ -412,28 +416,28 @@ class QNSPSAOptimizer:
         op_forward = self._get_operations(cost, args1, kwargs)
         op_inv = self._get_operations(cost, args2, kwargs)
 
-        new_ops = op_forward + [qml.adjoint(op) for op in reversed(op_inv)]
-        tape = qml.workflow.construct_tape(cost)(*args1, **kwargs)
-        return qml.tape.QuantumScript(new_ops, [qml.probs(wires=tape.wires.labels)])
+        new_ops = op_forward + [adjoint(op) for op in reversed(op_inv)]
+        tape = construct_tape(cost)(*args1, **kwargs)
+        return QuantumScript(new_ops, [probs(wires=tape.wires.labels)])
 
     @staticmethod
     def _get_operations(cost, args, kwargs):
-        tape = qml.workflow.construct_tape(cost)(*args, **kwargs)
+        tape = construct_tape(cost)(*args, **kwargs)
         return tape.operations
 
     def _apply_blocking(self, cost, args, kwargs, params_next):
-        tape = qml.workflow.construct_tape(cost)(*args, **kwargs)
+        tape = construct_tape(cost)(*args, **kwargs)
         tape_loss_curr = tape.copy(copy_operations=True)
 
         if not isinstance(params_next, list):
             params_next = [params_next]
 
-        tape = qml.workflow.construct_tape(cost)(*params_next, **kwargs)
+        tape = construct_tape(cost)(*params_next, **kwargs)
         tape_loss_next = tape.copy(copy_operations=True)
 
         program = cost.device.preprocess_transforms()
 
-        loss_curr, loss_next = qml.execute(
+        loss_curr, loss_next = execute(
             [tape_loss_curr, tape_loss_next], cost.device, None, transform_program=program
         )
 
diff --git a/pennylane/optimize/riemannian_gradient.py b/pennylane/optimize/riemannian_gradient.py
index 0bf03e594..76892f3d2 100644
--- a/pennylane/optimize/riemannian_gradient.py
+++ b/pennylane/optimize/riemannian_gradient.py
@@ -16,9 +16,15 @@ import warnings
 
 import numpy as np
 
-import pennylane as qml
+from pennylane import math
+from pennylane.ops import LinearCombination, PauliRot
+from pennylane.ops.functions import evolve
+from pennylane.pauli import pauli_group, pauli_word_to_string
 from pennylane.tape import QuantumScript, QuantumScriptBatch
+from pennylane.templates import ApproxTimeEvolution
+from pennylane.transforms.core import transform
 from pennylane.typing import PostprocessingFn
+from pennylane.workflow import QNode
 
 
 def null_postprocessing(results):
@@ -28,7 +34,7 @@ def null_postprocessing(results):
     return results[0]
 
 
-@qml.transform
+@transform
 def append_time_evolution(
     tape: QuantumScript, riemannian_gradient, t, n, exact=False
 ) -> tuple[QuantumScriptBatch, PostprocessingFn]:
@@ -71,16 +77,16 @@ def append_time_evolution(
 
     """
     if exact:
-        op = qml.evolve(riemannian_gradient, t)
+        op = evolve(riemannian_gradient, t)
     else:
-        op = qml.templates.ApproxTimeEvolution(riemannian_gradient, t, n)
+        op = ApproxTimeEvolution(riemannian_gradient, t, n)
 
     new_tape = tape.copy(operations=tape.operations + [op])
 
     return [new_tape], null_postprocessing
 
 
-@qml.transform
+@transform
 def algebra_commutator(tape, lie_algebra_basis_names, nqubits):
     """Calculate the Riemannian gradient in the Lie algebra with the parameter shift rule
     (see :meth:`RiemannianGradientOptimizer.get_omegas`).
@@ -104,17 +110,15 @@ def algebra_commutator(tape, lie_algebra_basis_names, nqubits):
     """
 
     wires = list(range(nqubits))
-    paulis_plus = (qml.PauliRot(np.pi / 2, pauli, wires=wires) for pauli in lie_algebra_basis_names)
-    paulis_minus = (
-        qml.PauliRot(-np.pi / 2, pauli, wires=wires) for pauli in lie_algebra_basis_names
-    )
+    paulis_plus = (PauliRot(np.pi / 2, pauli, wires=wires) for pauli in lie_algebra_basis_names)
+    paulis_minus = (PauliRot(-np.pi / 2, pauli, wires=wires) for pauli in lie_algebra_basis_names)
 
     tapes_plus = [tape.copy(operations=tape.operations + [op]) for op in paulis_plus]
     tapes_minus = [tape.copy(operations=tape.operations + [op]) for op in paulis_minus]
 
     def calculate_omegas(results):
-        results_plus = qml.math.hstack(results[: len(results) // 2])
-        results_minus = qml.math.hstack(results[len(results) // 2 :])
+        results_plus = math.hstack(results[: len(results) // 2])
+        results_minus = math.hstack(results[len(results) // 2 :])
         return results_plus - results_minus
 
     return tapes_plus + tapes_minus, calculate_omegas
@@ -225,7 +229,7 @@ class RiemannianGradientOptimizer:
     # pylint: disable=too-many-arguments, too-many-positional-arguments
     # pylint: disable=too-many-instance-attributes
     def __init__(self, circuit, stepsize=0.01, restriction=None, exact=False, trottersteps=1):
-        if not isinstance(circuit, qml.QNode):
+        if not isinstance(circuit, QNode):
             raise TypeError(f"circuit must be a QNode, received {type(circuit)}")
 
         self.circuit = circuit.update(diff_method=None)
@@ -235,7 +239,7 @@ class RiemannianGradientOptimizer:
         self.nqubits = len(circuit.device.wires)
 
         self.hamiltonian = circuit.func().obs
-        if not isinstance(self.hamiltonian, qml.ops.LinearCombination):
+        if not isinstance(self.hamiltonian, LinearCombination):
             raise TypeError(
                 f"circuit must return the expectation value of a Hamiltonian,"
                 f"received {type(circuit.func().obs)}"
@@ -248,7 +252,7 @@ class RiemannianGradientOptimizer:
                 f"optimizing a {self.nqubits} qubit circuit may be slow.",
                 UserWarning,
             )
-        if restriction is not None and not isinstance(restriction, qml.ops.LinearCombination):
+        if restriction is not None and not isinstance(restriction, LinearCombination):
             raise TypeError(f"restriction must be a Hamiltonian, received {type(restriction)}")
         (
             self.lie_algebra_basis_ops,
@@ -279,7 +283,7 @@ class RiemannianGradientOptimizer:
         non_zero_omegas = -omegas[omegas != 0]
         nonzero_idx = np.nonzero(omegas)[0]
 
-        lie_gradient = qml.Hamiltonian(
+        lie_gradient = LinearCombination(
             non_zero_omegas,
             [self.lie_algebra_basis_ops[i] for i in nonzero_idx],
         )
@@ -300,8 +304,8 @@ class RiemannianGradientOptimizer:
 
         # construct the corresponding pennylane observables
         wire_map = {i: i for i in range(self.nqubits)}
-        ops = list(restriction.ops) if restriction else list(qml.pauli.pauli_group(self.nqubits))
-        return ops, [qml.pauli.pauli_word_to_string(ps, wire_map=wire_map) for ps in ops]
+        ops = list(restriction.ops) if restriction else list(pauli_group(self.nqubits))
+        return ops, [pauli_word_to_string(ps, wire_map=wire_map) for ps in ops]
 
     def get_omegas(self):
         r"""Measure the coefficients of the Riemannian gradient with respect to a Pauli word basis.
diff --git a/pennylane/optimize/rotoselect.py b/pennylane/optimize/rotoselect.py
index c61b27ba6..e4aaba037 100644
--- a/pennylane/optimize/rotoselect.py
+++ b/pennylane/optimize/rotoselect.py
@@ -15,7 +15,7 @@
 
 import numpy as np
 
-import pennylane as qml
+from pennylane.ops import RX, RY, RZ
 
 from .qng import _flatten_np, _unflatten_np
 
@@ -93,7 +93,7 @@ class RotoselectOptimizer:
     # pylint: disable=too-few-public-methods
 
     def __init__(self, possible_generators=None):
-        self.possible_generators = possible_generators or [qml.RX, qml.RY, qml.RZ]
+        self.possible_generators = possible_generators or [RX, RY, RZ]
 
     def step_and_cost(self, objective_fn, x, generators, **kwargs):
         """Update trainable arguments with one step of the optimizer and return the corresponding
diff --git a/pennylane/optimize/rotosolve.py b/pennylane/optimize/rotosolve.py
index 88e2d4d68..6450cc65d 100644
--- a/pennylane/optimize/rotosolve.py
+++ b/pennylane/optimize/rotosolve.py
@@ -19,7 +19,9 @@ from inspect import signature
 import numpy as np
 import scipy as sp
 
-import pennylane as qml
+from pennylane import math
+from pennylane.fourier import reconstruct  # tach-ignore
+from pennylane.workflow import QNode
 
 
 def _brute_optimizer(fun, num_steps, bounds=None, **kwargs):
@@ -64,7 +66,7 @@ def _validate_inputs(requires_grad, args, nums_frequency, spectra):
             _spectra = spectra.get(arg_name, {})
             all_keys = set(_nums_frequency) | set(_spectra)
 
-            shape = qml.math.shape(arg)
+            shape = math.shape(arg)
             indices = np.ndindex(shape) if len(shape) > 0 else [()]
             for par_idx in indices:
                 if par_idx not in all_keys:
@@ -92,11 +94,11 @@ def _restrict_to_univariate(fn, arg_idx, par_idx, args, kwargs):
         function is added to the marked parameter.
     """
     the_arg = args[arg_idx]
-    if len(qml.math.shape(the_arg)) == 0:
-        shift_vec = qml.math.ones_like(the_arg)
+    if len(math.shape(the_arg)) == 0:
+        shift_vec = math.ones_like(the_arg)
     else:
-        shift_vec = qml.math.zeros_like(the_arg)
-        shift_vec = qml.math.scatter_element_add(shift_vec, par_idx, 1.0)
+        shift_vec = math.zeros_like(the_arg)
+        shift_vec = math.scatter_element_add(shift_vec, par_idx, 1.0)
 
     def _univariate_fn(x):
         return fn(*args[:arg_idx], the_arg + shift_vec * x, *args[arg_idx + 1 :], **kwargs)
@@ -419,10 +421,10 @@ class RotosolveOptimizer:
 
         """
         # todo: does this signature call cover all cases?
-        sign_fn = objective_fn.func if isinstance(objective_fn, qml.QNode) else objective_fn
+        sign_fn = objective_fn.func if isinstance(objective_fn, QNode) else objective_fn
         arg_names = list(signature(sign_fn).parameters.keys())
         requires_grad = {
-            arg_name: qml.math.requires_grad(arg) for arg_name, arg in zip(arg_names, args)
+            arg_name: math.requires_grad(arg) for arg_name, arg in zip(arg_names, args)
         }
         nums_frequency = nums_frequency or {}
         spectra = spectra or {}
@@ -446,7 +448,7 @@ class RotosolveOptimizer:
             if not requires_grad[arg_name]:
                 before_args.append(arg)
                 continue
-            shape = qml.math.shape(arg)
+            shape = math.shape(arg)
             indices = np.ndindex(shape) if len(shape) > 0 else [()]
             for par_idx in indices:
                 _fun_at_zero = fun_at_zero if first_substep_in_step else None
@@ -463,7 +465,7 @@ class RotosolveOptimizer:
                     )
                     freq = 1.0 if num_freq is not None else spectrum[spectrum > 0][0]
                     x_min, y_min = self.min_analytic(univariate, freq, _fun_at_zero)
-                    arg = qml.math.scatter_element_add(arg, par_idx, x_min)
+                    arg = math.scatter_element_add(arg, par_idx, x_min)
 
                 else:
                     ids = {arg_name: (par_idx,)}
@@ -473,9 +475,7 @@ class RotosolveOptimizer:
                     _spectra = {arg_name: {par_idx: spectrum}} if spectrum is not None else None
 
                     # Set up the reconstruction function
-                    recon_fn = qml.fourier.reconstruct(
-                        objective_fn, ids, _nums_frequency, _spectra, shifts
-                    )
+                    recon_fn = reconstruct(objective_fn, ids, _nums_frequency, _spectra, shifts)
                     # Perform the reconstruction
                     recon = recon_fn(*before_args, arg, *after_args, f0=_fun_at_zero, **kwargs)[
                         arg_name
@@ -485,7 +485,7 @@ class RotosolveOptimizer:
                     x_min, y_min = self._min_numeric(recon, spectrum)
 
                     # Update the currently treated argument
-                    arg = qml.math.scatter_element_add(arg, par_idx, x_min - arg[par_idx])
+                    arg = math.scatter_element_add(arg, par_idx, x_min - arg[par_idx])
                 first_substep_in_step = False
 
                 if full_output:
@@ -612,8 +612,8 @@ class RotosolveOptimizer:
         """
         opt_kwargs = self.substep_kwargs.copy()
         if "bounds" not in self.substep_kwargs:
-            spectrum = qml.math.array(spectrum)
-            half_width = np.pi / qml.math.min(spectrum[spectrum > 0])
+            spectrum = math.array(spectrum)
+            half_width = np.pi / math.min(spectrum[spectrum > 0])
             opt_kwargs["bounds"] = ((-half_width, half_width),)
 
         x_min, y_min = self.substep_optimizer(objective_fn, **opt_kwargs)
diff --git a/pennylane/optimize/shot_adaptive.py b/pennylane/optimize/shot_adaptive.py
index 87b33943d..7f855f795 100644
--- a/pennylane/optimize/shot_adaptive.py
+++ b/pennylane/optimize/shot_adaptive.py
@@ -18,7 +18,12 @@ from copy import copy
 import numpy as np
 import scipy as sp
 
-import pennylane as qml
+from pennylane import math, measurements
+from pennylane._grad import jacobian
+from pennylane.ops import LinearCombination
+from pennylane.queuing import apply
+from pennylane.tape import make_qscript
+from pennylane.workflow import QNode, construct_tape
 
 from .gradient_descent import GradientDescentOptimizer
 
@@ -174,6 +179,7 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
           <https://quantum-journal.org/papers/q-2020-05-11-263/>`__ (2020).
     """
 
+    # pylint: disable = too-many-positional-arguments
     def __init__(self, min_shots, term_sampling=None, mu=0.99, b=1e-6, stepsize=0.07):
         self.term_sampling = term_sampling
         self.trainable_args = set()
@@ -243,10 +249,10 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
                 continue
 
             def func(*qnode_args, **qnode_kwargs):
-                qs = qml.tape.make_qscript(base_func)(*qnode_args, **qnode_kwargs)
+                qs = make_qscript(base_func)(*qnode_args, **qnode_kwargs)
                 for op in qs.operations:
-                    qml.apply(op)
-                return qml.expval(o)  # pylint:disable=cell-var-from-loop
+                    apply(op)
+                return measurements.expval(o)  # pylint:disable=cell-var-from-loop
 
             qnode.func = func
 
@@ -256,12 +262,12 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
 
                 def cost(*args, **kwargs):
                     # pylint: disable=cell-var-from-loop
-                    return qml.math.stack(qnode(*args, **kwargs))
+                    return math.stack(qnode(*args, **kwargs))
 
             else:
                 cost = qnode
 
-            jacs = qml.jacobian(cost, argnum=argnums)(*args, **kwargs, shots=new_shots)
+            jacs = jacobian(cost, argnum=argnums)(*args, **kwargs, shots=new_shots)
 
             if s == 1:
                 jacs = [np.expand_dims(j, 0) for j in jacs]
@@ -308,11 +314,11 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
         """Compute the single shot gradients of a QNode."""
         self.check_device(qnode.device)
 
-        tape = qml.workflow.construct_tape(qnode)(*args, **kwargs)
+        tape = construct_tape(qnode)(*args, **kwargs)
         [expval] = tape.measurements
         coeffs, observables = (
             expval.obs.terms()
-            if isinstance(expval.obs, qml.ops.LinearCombination)
+            if isinstance(expval.obs, LinearCombination)
             else ([1.0], [expval.obs])
         )
 
@@ -333,9 +339,9 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
         new_shots = [(1, int(self.max_shots))]
 
         def cost(*args, **kwargs):
-            return qml.math.stack(qnode(*args, **kwargs, shots=new_shots))
+            return math.stack(qnode(*args, **kwargs, shots=new_shots))
 
-        grads = [qml.jacobian(cost, argnum=i)(*args, **kwargs) for i in self.trainable_args]
+        grads = [jacobian(cost, argnum=i)(*args, **kwargs) for i in self.trainable_args]
 
         return grads
 
@@ -354,7 +360,7 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
             tuple[array[float], array[float]]: a tuple of NumPy arrays containing the gradient
             :math:`\nabla f(x^{(t)})` and the variance of the gradient
         """
-        if isinstance(objective_fn, qml.QNode) or hasattr(objective_fn, "device"):
+        if isinstance(objective_fn, QNode) or hasattr(objective_fn, "device"):
             grads = self._single_shot_qnode_gradients(objective_fn, args, kwargs)
         else:
             raise ValueError(
diff --git a/tach.toml b/tach.toml
index 421448059..73776fed3 100644
--- a/tach.toml
+++ b/tach.toml
@@ -142,6 +142,11 @@ path = "pennylane.shadows"
 layer = "auxiliary"
 cannot_depend_on = ["pennylane.ftqc", "pennylane.labs"]
 
+[[modules]]
+path = "pennylane.optimize"
+layer = "auxiliary"
+cannot_depend_on = ["pennylane.ftqc", "pennylane.labs"]
+
 ## TERTIARY modules
 # These modules can only import from CORE or AUXILIARY modules.
 
@@ -255,10 +260,6 @@ cannot_depend_on = ["pennylane.ftqc", "pennylane.labs"]
 path = "pennylane.qnn"
 cannot_depend_on = ["pennylane.ftqc", "pennylane.labs"]
 
-[[modules]]
-path = "pennylane.optimize"
-cannot_depend_on = ["pennylane.ftqc", "pennylane.labs"]
-
 [[modules]]
 path = "pennylane.noise"
 cannot_depend_on = ["pennylane.ftqc", "pennylane.labs"]


Test patch:
diff --git a/tests/optimize/test_optimize_shot_adaptive.py b/tests/optimize/test_optimize_shot_adaptive.py
index ed497bb1c..266f5f380 100644
--- a/tests/optimize/test_optimize_shot_adaptive.py
+++ b/tests/optimize/test_optimize_shot_adaptive.py
@@ -17,6 +17,7 @@ import pytest
 
 import pennylane as qml
 from pennylane import numpy as np
+from pennylane import optimize
 
 
 class TestExceptions:
@@ -556,7 +557,8 @@ class TestQNodeWeightedRandomSampling:
         weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))
 
         opt = qml.ShotAdaptiveOptimizer(min_shots=10)
-        spy = mocker.spy(qml, "jacobian")
+
+        spy = mocker.spy(optimize.shot_adaptive, "jacobian")
         mocker.patch(
             "scipy.stats._multivariate.multinomial_gen.rvs", return_value=np.array([[4, 0, 6]])
         )
@@ -582,7 +584,7 @@ class TestQNodeWeightedRandomSampling:
 
         opt = qml.ShotAdaptiveOptimizer(min_shots=10)
 
-        spy = mocker.spy(qml, "jacobian")
+        spy = mocker.spy(optimize.shot_adaptive, "jacobian")
         mocker.patch(
             "scipy.stats._multivariate.multinomial_gen.rvs", return_value=np.array([[4, 1, 5]])
         )