Repository: PennyLaneAI/pennylane
Pull request #7452 (2025-05-20 14:46:56)

Implement new QNG Jax-based optimizer compatible with (q)jit
For PL's non-quantum optimizers (e.g. `GradientDescentOptimizer`, `AdamOptimizer`, etc.), the user is currently redirected to leverage `optax`-native functionality (see [here](https://optax.readthedocs.io/en/latest/api/optimizers.html)) instead of `qml.*` when he wants to enable `(q)jit` compilation. However, all the quantum-specific optimizers available in PL (e.g. `QNGOptimizer`, `AdaptiveOptimizer`, etc.) are not implemented at all in `optax`. The goal is to implement a new Quantum Natural Gradient (QNG) optimizer version compatible with `jax.jit` and `qml.qjit`.

Here is a minimal example of what we are aiming for:
```python
import pennylane as qml
import jax
import jax.numpy as jnp
from functools import partial

dev = qml.device("default.qubit", wires=1)

@qml.qnode(dev)
def circuit(params):
    qml.RX(params[0], wires=0)
    qml.RY(params[1], wires=0)
    return qml.expval(qml.Z(0))

opt = qml.QNGOptimizerQJIT(stepsize=0.1)
params = jnp.array([0.1, 0.2])

step = jax.jit(partial(opt.step, circuit))

state = opt.init(params)
for _ in range(1000):
   params, state = step(params, state)
```

This likely requires some preliminary work to make `qml.metric_tensor` Catalyst compatible (@albi3ro is taking care of this).

[sc-91241]

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.

Patch:
diff --git a/doc/introduction/interfaces.rst b/doc/introduction/interfaces.rst
index 4e9117290..257149bd5 100644
--- a/doc/introduction/interfaces.rst
+++ b/doc/introduction/interfaces.rst
@@ -186,6 +186,18 @@ Check out the `JAXopt <https://github.com/google/jaxopt>`_ and the `Optax
 <https://optax.readthedocs.io/en/latest/>`_ packages to find optimizers for the
 :ref:`PennyLane JAX framework <jax_interf>`.
 
+For quantum-specific optimizers (not available in `Optax <https://optax.readthedocs.io/en/latest/>`_),
+PennyLane offers a built-in Jax-based implementation compatible with ``jax.jit`` and ``qml.qjit``.
+
+:html:`<div class="summary-table">`
+
+.. autosummary::
+    :nosignatures:
+
+    ~pennylane.QNGOptimizerQJIT
+
+:html:`</div>`
+
 Gradients
 ---------
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 4d7e95d40..371547da9 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,42 @@
 
 <h3>New features since last release</h3>
 
+* Leveraging quantum just-in-time compilation to optimize parameterized hybrid workflows with the quantum 
+  natural gradient optimizer is now possible with the new :class:`~.QNGOptimizerQJIT` optimizer. 
+  [(#7452)](https://github.com/PennyLaneAI/pennylane/pull/7452)
+  
+  The :class:`~.QNGOptimizerQJIT` optimizer offers a `jax.jit`- and `qml.qjit`-compatible analogue to the existing 
+  :class:`~.QNGOptimizer` with an Optax-like interface:
+
+  ```python
+  import pennylane as qml
+  import jax.numpy as jnp
+
+  @qml.qjit(autograph=True)
+  def workflow():
+      dev = qml.device("lightning.qubit", wires=2)
+  
+      @qml.qnode(dev)
+      def circuit(params):
+          qml.RX(params[0], wires=0)
+          qml.RY(params[1], wires=1)
+          return qml.expval(qml.Z(0) + qml.X(1))
+  
+      opt = qml.QNGOptimizerQJIT(stepsize=0.2)
+  
+      params = jnp.array([0.1, 0.2])
+      state = opt.init(params)
+      for _ in range(100):
+          params, state = opt.step(circuit, params, state)
+  
+      return params
+  ```
+
+  ```pycon
+  >>> workflow()
+  Array([ 3.14159265, -1.57079633], dtype=float64)
+  ```
+
 <h4>State-of-the-art templates and decompositions 🐝</h4>
 
 * The decompositions of `SingleExcitation`, `SingleExcitationMinus` and `SingleExcitationPlus`
diff --git a/pennylane/optimize/__init__.py b/pennylane/optimize/__init__.py
index ec2d75d7f..8ffe29ced 100644
--- a/pennylane/optimize/__init__.py
+++ b/pennylane/optimize/__init__.py
@@ -32,6 +32,7 @@ from .rotoselect import RotoselectOptimizer
 from .rotosolve import RotosolveOptimizer
 from .shot_adaptive import ShotAdaptiveOptimizer
 from .spsa import SPSAOptimizer
+from .qng_qjit import QNGOptimizerQJIT
 
 # Optimizers to display in the docs
 __all__ = [
@@ -50,4 +51,5 @@ __all__ = [
     "RotoselectOptimizer",
     "ShotAdaptiveOptimizer",
     "SPSAOptimizer",
+    "QNGOptimizerQJIT",
 ]
diff --git a/pennylane/optimize/qng_qjit.py b/pennylane/optimize/qng_qjit.py
new file mode 100644
index 000000000..b571692d3
--- /dev/null
+++ b/pennylane/optimize/qng_qjit.py
@@ -0,0 +1,272 @@
+# Copyright 2018-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.
+"""Quantum natural gradient optimizer for Jax/Catalyst interface"""
+
+from pennylane import math
+from pennylane.compiler import active_compiler
+from pennylane.gradients.metric_tensor import metric_tensor
+from pennylane.workflow import QNode
+
+has_catalyst = True
+try:
+    import catalyst
+except ModuleNotFoundError:
+    has_catalyst = False
+
+has_jax = True
+try:
+    import jax
+except ModuleNotFoundError:
+    has_jax = False
+
+
+class QNGOptimizerQJIT:
+    r"""Optax-like and ``jax.jit``/``qml.qjit``-compatible implementation of the ``QNGOptimizer``,
+    a step- and parameter-dependent learning rate optimizer, leveraging a reparameterization of
+    the optimization space based on the Fubini-Study metric tensor.
+
+    For more theoretical details, see the :class:`~.QNGOptimizer` documentation.
+
+    .. note::
+
+        Please be aware of the following:
+
+            - As with ``QNGOptimizer``, ``QNGOptimizerQJIT`` supports a single QNode to encode the objective function.
+
+            - ``QNGOptimizerQJIT`` does not support any QNode with multiple arguments. A potential workaround
+              would be to combine all parameters into a single objective function argument.
+
+            - ``QNGOptimizerQJIT`` does not work correctly if there is any classical processing in the QNode circuit
+              (e.g., ``2 * theta`` as a gate parameter).
+
+    **Example:**
+
+    Consider a hybrid workflow to optimize an objective function defined by a quantum circuit.
+    To make the optimization faster, the entire workflow can be just-in-time compiled using
+    the ``qml.qjit`` decorator:
+
+    .. code-block:: python
+
+        import pennylane as qml
+        import jax.numpy as jnp
+
+        @qml.qjit(autograph=True)
+        def workflow():
+            dev = qml.device("lightning.qubit", wires=2)
+
+            @qml.qnode(dev)
+            def circuit(params):
+                qml.RX(params[0], wires=0)
+                qml.RY(params[1], wires=1)
+                return qml.expval(qml.Z(0) + qml.X(1))
+
+            opt = qml.QNGOptimizerQJIT(stepsize=0.2)
+
+            params = jnp.array([0.1, 0.2])
+            state = opt.init(params)
+            for _ in range(100):
+                params, state = opt.step(circuit, params, state)
+
+            return params
+
+    >>> workflow()
+    Array([ 3.14159265, -1.57079633], dtype=float64)
+
+    Make sure you are using the ``lightning.qubit`` device along with ``qml.qjit`` with ``autograph`` enabled.
+
+    Using the ``jax.jit`` decorator for the entire workflow is not recommended since it
+    may lead to a significative compilation time and no runtime benefits.
+    However, ``jax.jit`` can be used with the ``default.qubit`` device to just-in-time
+    compile the ``step`` (or ``step_and_cost``) method of the optimizer.
+    For example:
+
+    .. code-block:: python
+
+        import pennylane as qml
+        import jax.numpy as jnp
+        import jax
+        from functools import partial
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(params):
+            qml.RX(params[0], wires=0)
+            qml.RY(params[1], wires=1)
+            return qml.expval(qml.Z(0) + qml.X(1))
+
+        opt = qml.QNGOptimizerQJIT(stepsize=0.2)
+        step = jax.jit(partial(opt.step, circuit))
+
+        params = jnp.array([0.1, 0.2])
+        state = opt.init(params)
+        for _ in range(100):
+            params, state = step(params, state)
+
+    >>> params
+    Array([ 3.14159265, -1.57079633], dtype=float64)
+
+    Keyword Args:
+        stepsize=0.01 (float): the user-defined stepsize hyperparameter
+        approx="block-diag" (str): approximation method for the metric tensor.
+
+            - If ``None``, the full metric tensor is computed
+
+            - If ``"block-diag"``, the block-diagonal approximation is computed, reducing
+              the number of evaluated circuits significantly
+
+            - If ``"diag"``, the diagonal approximation is computed, slightly
+              reducing the classical overhead but not the quantum resources
+              (compared to ``"block-diag"``)
+
+        lam=0 (float): metric tensor regularization to be applied at each optimization step
+    """
+
+    def __init__(self, stepsize=0.01, approx="block-diag", lam=0):
+        self.stepsize = stepsize
+        self.approx = approx
+        self.lam = lam
+
+    def init(self, params):
+        """Return the initial state of the optimizer.
+
+        Args:
+            params (array): QNode parameters
+
+        Returns:
+            None
+
+        .. note::
+
+            Since the Quantum Natural Gradient (QNG) algorithm doesn't actually require any particular state,
+            this method always returns an empty ``None`` state. However, it is provided to match
+            the ``optax``-like interface for all Jax-based quantum-specific optimizers.
+
+        """
+        # pylint:disable=unused-argument
+        # pylint:disable=no-self-use
+        return None
+
+    def step(self, qnode, params, state, **kwargs):
+        """Update the QNode parameters and the optimizer's state for a single optimization step.
+
+        Args:
+            qnode (QNode): QNode objective function to be optimized
+            params (array): QNode parameters to be updated
+            state: current state of the optimizer
+            **kwargs : variable-length keyword arguments for the QNode
+
+        Returns:
+            tuple: (new parameters values, new optimizer's state)
+
+        .. note::
+
+            Since the Quantum Natural Gradient (QNG) algorithm doesn't actually require any particular state,
+            the ``state`` object is never really updated in this case. However, it is carried over the
+            optimization to match the ``optax``-like interface for all Jax-based quantum-specific optimizers.
+
+        """
+        mt = self._get_metric_tensor(qnode, params, **kwargs)
+        grad = self._get_grad(qnode, params, **kwargs)
+        new_params, new_state = self._apply_grad(mt, grad, params, state)
+        return new_params, new_state
+
+    def step_and_cost(self, qnode, params, state, **kwargs):
+        """Update the QNode parameters and the optimizer's state for a single optimization step
+        and return the corresponding objective function value prior to the step.
+
+        Args:
+            qnode (QNode): QNode objective function to be optimized
+            params (array): QNode parameters to be updated
+            state: current state of the optimizer
+            **kwargs : variable-length keyword arguments for the QNode
+
+        Returns:
+            tuple: (new parameters values, new optimizer's state, objective function value)
+
+        .. note::
+
+            Since the Quantum Natural Gradient (QNG) algorithm doesn't actually require any particular state,
+            the ``state`` object is never really updated in this case. However, it is carried over the
+            optimization to match the ``optax``-like interface for all Jax-based quantum-specific optimizers.
+
+        """
+        mt = self._get_metric_tensor(qnode, params, **kwargs)
+        cost, grad = self._get_value_and_grad(qnode, params, **kwargs)
+        new_params, new_state = self._apply_grad(mt, grad, params, state)
+        return new_params, new_state, cost
+
+    @staticmethod
+    def _get_grad(qnode, params, **kwargs):
+        """Return the gradient of the QNode objective function at the given point. The method is implemented to dispatch
+        to Catalyst when it is required (e.g. when using ``qml.qjit``) or to fall back to Jax otherwise.
+
+        Raise an ``ModuleNotFoundError`` if the required package is not installed.
+        """
+        if active_compiler() == "catalyst":
+            if has_catalyst:
+                return catalyst.grad(qnode)(params, **kwargs)
+            raise ModuleNotFoundError("Catalyst is required.")  # pragma: no cover
+        if has_jax:
+            return jax.grad(qnode)(params, **kwargs)
+        raise ModuleNotFoundError("Jax is required.")  # pragma: no cover
+
+    @staticmethod
+    def _get_value_and_grad(qnode, params, **kwargs):
+        """Return the value and the gradient of the QNode objective function at the given point. The method is implemented
+        to dispatch to Catalyst when it is required (e.g. when using ``qml.qjit``) or to fall back to Jax otherwise.
+
+        Raise an ``ModuleNotFoundError`` if the required package is not installed.
+        """
+        if active_compiler() == "catalyst":
+            if has_catalyst:
+                return catalyst.value_and_grad(qnode)(params, **kwargs)
+            raise ModuleNotFoundError("Catalyst is required.")  # pragma: no cover
+        if has_jax:
+            return jax.value_and_grad(qnode)(params, **kwargs)
+        raise ModuleNotFoundError("Jax is required.")  # pragma: no cover
+
+    def _get_metric_tensor(self, qnode, params, **kwargs):
+        """Compute the metric tensor of the QNode objective function at the given point using the method specified
+        by the optimizer's ``approx`` attribute. It returns the reshaped matrix after applying the regularization
+        given by the optimizer's ``lam`` attribute.
+
+        Raise a ``ValueError`` if the given objective function is not encoded as a QNode.
+        """
+        # pylint: disable=not-callable
+        if not isinstance(qnode, QNode):
+            raise ValueError(
+                "The objective function must be encoded as a single QNode to use the Quantum Natural Gradient optimizer."
+            )
+        mt = metric_tensor(qnode, approx=self.approx)(params, **kwargs)
+        # reshape tensor into a matrix (acting on the flat grad vector)
+        shape = math.shape(mt)
+        size = 1 if shape == () else math.prod(shape[: len(shape) // 2])
+        mt_matrix = math.reshape(mt, (size, size))
+        # apply regularization for matrix inversion
+        if self.lam != 0:
+            mt_matrix += self.lam * math.eye(size, like=mt_matrix)
+        return mt_matrix
+
+    def _apply_grad(self, mt, grad, params, state):
+        """Update the parameter array ``params`` for a single optimization step according to the Quantum
+        Natural Gradient algorithm. The method doesn't perform any transformation on ``state`` since the QNG
+        optimizer doesn't actually require any particular state.
+        """
+        shape = math.shape(grad)
+        grad_flat = math.flatten(grad)
+        update_flat = math.linalg.pinv(mt) @ grad_flat
+        update = math.reshape(update_flat, shape)
+        new_params = params - self.stepsize * update
+        return new_params, state


Test patch:
diff --git a/tests/optimize/test_qng_qjit.py b/tests/optimize/test_qng_qjit.py
new file mode 100644
index 000000000..76025d9c6
--- /dev/null
+++ b/tests/optimize/test_qng_qjit.py
@@ -0,0 +1,418 @@
+# Copyright 2018-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.
+"""Test Jax-based Catalyst-compatible QNG optimizer"""
+
+from functools import partial
+
+import numpy as np
+import pytest
+
+import pennylane as qml
+
+dev_names = (
+    "default.qubit",
+    "lightning.qubit",
+)
+
+
+def circuit(params):
+    """Simple circuit to use for testing."""
+    qml.RX(params[0], wires=0)
+    qml.RY(params[1], wires=1)
+    return qml.expval(qml.Z(wires=0))
+
+
+class TestBasics:
+    """Test basic properties of the QNGOptimizerQJIT."""
+
+    def test_initialization_default(self):
+        """Test that initializing QNGOptimizerQJIT with default values works."""
+        opt = qml.QNGOptimizerQJIT()
+        assert opt.stepsize == 0.01
+        assert opt.approx == "block-diag"
+        assert opt.lam == 0
+
+    def test_initialization_custom(self):
+        """Test that initializing QNGOptimizerQJIT with custom values works."""
+        opt = qml.QNGOptimizerQJIT(stepsize=0.05, approx="diag", lam=1e-9)
+        assert opt.stepsize == 0.05
+        assert opt.approx == "diag"
+        assert opt.lam == 1e-9
+
+    def test_init_none_state(self):
+        """Test that the QNGOptimizerQJIT state is initialized to `None`."""
+        opt = qml.QNGOptimizerQJIT()
+        state = opt.init([0.1, 0.2])
+        assert state is None
+
+
+class TestGradients:
+    """Test gradient computation in the QNGOptimizerQJIT."""
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", dev_names)
+    def test_get_grad_jax(self, dev_name):
+        """Test `_get_grad` method with Jax interface."""
+        # pylint:disable=protected-access
+        import jax.numpy as jnp
+
+        device = qml.device(dev_name, wires=2)
+        qnode = qml.QNode(circuit, device=device)
+
+        opt = qml.QNGOptimizerQJIT()
+        params = [0.1, 0.2]
+
+        params_jax = jnp.array(params)
+        grad_jax = opt._get_grad(qnode, params_jax)
+
+        grad_exact = jnp.array([-np.sin(params[0]), 0.0])
+
+        assert np.allclose(grad_exact, grad_jax)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", dev_names)
+    def test_get_value_and_grad_jax(self, dev_name):
+        """Test `_get_value_and_grad` method with Jax interface."""
+        # pylint:disable=protected-access
+        import jax.numpy as jnp
+
+        device = qml.device(dev_name, wires=2)
+        qnode = qml.QNode(circuit, device=device)
+
+        opt = qml.QNGOptimizerQJIT()
+        params = [0.1, 0.2]
+
+        params_jax = jnp.array(params)
+        cost_jax, grad_jax = opt._get_value_and_grad(qnode, params_jax)
+
+        cost_exact = qnode(params)
+        grad_exact = jnp.array([-np.sin(params[0]), 0.0])
+
+        assert np.allclose(cost_exact, cost_jax)
+        assert np.allclose(grad_exact, grad_jax)
+
+
+class TestMetricTensor:
+    """Test that the metric tensor computation works with `approx` and `lam`."""
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", dev_names)
+    def test_no_approx(self, dev_name):
+        """Test that the full metric tensor is computed for `approx=None`."""
+        # pylint:disable=protected-access
+        import jax.numpy as jnp
+
+        @qml.qnode(qml.device(dev_name))
+        def circ(params):
+            qml.RY(eta, wires=0)
+            qml.RX(params[0], wires=0)
+            qml.RY(params[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        eta = 0.7
+        params = jnp.array([0.11, 0.412])
+
+        opt = qml.QNGOptimizerQJIT(approx=None)
+        mt = opt._get_metric_tensor(circ, params)
+
+        # computing the expected metric tensor requires some manual calculation
+        x = params[0]
+        first_term = np.eye(2) / 4
+        vec_potential = np.array([-0.5j * np.sin(eta), 0.5j * np.sin(x) * np.cos(eta)])
+        second_term = np.real(np.outer(vec_potential.conj(), vec_potential))
+        expected_mt = first_term - second_term
+
+        assert np.allclose(mt, expected_mt)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", dev_names)
+    def test_with_approx(self, dev_name):
+        """Test that the approximated metric tensor is computed for `approx=block-diag` and `approx=diag`."""
+        # pylint:disable=protected-access
+        import jax.numpy as jnp
+
+        @qml.qnode(qml.device(dev_name))
+        def circ(params):
+            qml.RY(eta, wires=0)
+            qml.RX(params[0], wires=0)
+            qml.RY(params[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        eta = 0.7
+        params = jnp.array([0.11, 0.412])
+
+        blockdiag_opt = qml.QNGOptimizerQJIT(approx="block-diag")
+        blockdiag_mt = blockdiag_opt._get_metric_tensor(circ, params)
+
+        diag_opt = qml.QNGOptimizerQJIT(approx="diag")
+        diag_mt = diag_opt._get_metric_tensor(circ, params)
+
+        # computing the expected metric tensor requires some manual calculation
+        x = params[0]
+        first_term = np.eye(2) / 4
+        vec_potential = np.array([-0.5j * np.sin(eta), 0.5j * np.sin(x) * np.cos(eta)])
+        second_term = np.real(np.outer(vec_potential.conj(), vec_potential))
+        expected_mt = np.diag(np.diag(first_term - second_term))
+
+        assert np.allclose(blockdiag_mt, expected_mt)
+        assert np.allclose(diag_mt, expected_mt)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", dev_names)
+    def test_lam(self, dev_name):
+        """Test that the regularization `lam` is used correctly."""
+        # pylint:disable=protected-access
+        import jax.numpy as jnp
+
+        @qml.qnode(qml.device(dev_name))
+        def circ(params):
+            qml.RY(eta, wires=0)
+            qml.RX(params[0], wires=0)
+            qml.RY(params[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        eta = np.pi
+        params = jnp.array([eta / 2, 0.412])
+
+        stepsize = 1.0
+        lam = 1e-11
+
+        opt_with_lam = qml.QNGOptimizerQJIT(stepsize=stepsize, approx=None, lam=lam)
+        state_with_lam = opt_with_lam.init(params)
+        new_params_with_lam, _ = opt_with_lam.step(circ, params, state_with_lam)
+        mt_with_lam = opt_with_lam._get_metric_tensor(circ, params)
+
+        opt = qml.QNGOptimizerQJIT(stepsize=stepsize, approx=None)
+        state = opt.init(params)
+        new_params, _ = opt.step(circ, params, state)
+        mt = opt._get_metric_tensor(circ, params)
+
+        # computing the expected metric tensor requires some manual calculation
+        x, y = params
+        first_term = np.eye(2) / 4
+        vec_potential = np.array([-0.5j * np.sin(eta), 0.5j * np.sin(x) * np.cos(eta)])
+        second_term = np.real(np.outer(vec_potential.conj(), vec_potential))
+        expected_mt = first_term - second_term
+
+        assert np.allclose(mt_with_lam, expected_mt + np.eye(2) * lam)
+        assert np.allclose(mt, expected_mt)
+
+        # with regularization y gets updated, without regularization it does not
+        assert not np.isclose(new_params_with_lam[1], y)
+        assert np.isclose(new_params[1], y)
+
+
+# pylint:disable=too-few-public-methods
+class TestExceptions:
+    """Test exceptions are raised for incorrect usage."""
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", dev_names)
+    def test_obj_func_not_a_qnode(self, dev_name):
+        """Test that if the objective function is not a QNode, an error is raised."""
+        import jax.numpy as jnp
+
+        @qml.qnode(qml.device(dev_name))
+        def circ(a):
+            qml.RX(a, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        def cost(a):
+            return circ(a)
+
+        opt = qml.QNGOptimizerQJIT()
+        params = jnp.array(0.5)
+        state = opt.init(params)
+
+        with pytest.raises(
+            ValueError,
+            match="The objective function must be encoded as a single QNode to use the Quantum Natural Gradient optimizer.",
+        ):
+            opt.step(cost, params, state)
+
+
+class TestOptimize:
+    """Test basic optimization integration."""
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", dev_names)
+    def test_step_and_cost(self, dev_name):
+        """Test that the step and step_and_cost methods are returning the correct result."""
+        import jax.numpy as jnp
+
+        @qml.qnode(qml.device(dev_name))
+        def circ(params):
+            qml.RX(params[0], wires=0)
+            qml.RY(params[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        params = [0.31, 0.842]
+        params_qml = qml.numpy.array(params)
+        params_jax = jnp.array(params)
+
+        opt = qml.QNGOptimizerQJIT(stepsize=0.05)
+        state = opt.init(params_jax)
+
+        new_params1, _ = opt.step(circ, params_jax, state)
+        new_params2, _, cost = opt.step_and_cost(circ, params_jax, state)
+
+        expected_mt = np.array([0.25, (np.cos(params[0]) ** 2) / 4])
+        expected_params = params_qml - opt.stepsize * qml.grad(circ)(params_qml) / expected_mt
+        expected_cost = circ(params)
+
+        assert np.allclose(new_params1, expected_params)
+        assert np.allclose(new_params2, expected_params)
+        assert np.allclose(cost, expected_cost)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", dev_names)
+    def test_step_and_cost_with_gen_hamiltonian(self, dev_name):
+        """Test that the step and step_and_cost methods are returning the correct result
+        when the generator of an operator is a Hamiltonian."""
+        import jax.numpy as jnp
+
+        @qml.qnode(qml.device(dev_name))
+        def circ(params):
+            qml.DoubleExcitation(params[0], wires=[0, 1, 2, 3])
+            qml.RY(params[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        params = [0.31, 0.842]
+        params_qml = qml.numpy.array(params)
+        params_jax = jnp.array(params)
+
+        opt = qml.QNGOptimizerQJIT(stepsize=0.05)
+        state = opt.init(params_jax)
+
+        new_params1, _ = opt.step(circ, params_jax, state)
+        new_params2, _, cost = opt.step_and_cost(circ, params_jax, state)
+
+        expected_cost = circ(params)
+        expected_mt = np.array([1 / 16, 1 / 4])
+        expected_params = params_qml - opt.stepsize * qml.grad(circ)(params_qml) / expected_mt
+
+        assert np.allclose(cost, expected_cost)
+        assert np.allclose(new_params1, expected_params)
+        assert np.allclose(new_params2, expected_params)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", dev_names)
+    def test_qubit_rotations_circuit(self, tol, dev_name):
+        """Test that a simple qubit rotations circuit gets optimized correctly, checking params and cost at each step."""
+        import jax.numpy as jnp
+
+        @qml.qnode(qml.device(dev_name))
+        def circ(params):
+            qml.RX(params[0], wires=0)
+            qml.RY(params[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        def grad(params):
+            """Returns the gradient of the above circuit."""
+            da = -np.sin(params[0]) * np.cos(params[1])
+            db = -np.cos(params[0]) * np.sin(params[1])
+            return np.array([da, db])
+
+        opt = qml.QNGOptimizerQJIT(stepsize=0.1)
+        params = jnp.array([0.011, 0.012])
+        expected_params = jnp.array([0.011, 0.012])
+        state = opt.init(params)
+
+        num_steps = 30
+        for _ in range(num_steps):
+            params, state, cost = opt.step_and_cost(circ, params, state)
+
+            expected_cost = circ(expected_params)
+            expected_mt = np.array([0.25, (np.cos(expected_params[0]) ** 2) / 4])
+            expected_params -= opt.stepsize * grad(expected_params) / expected_mt
+
+            assert np.allclose(cost, expected_cost, atol=tol, rtol=0)
+            assert np.allclose(params, expected_params, atol=tol, rtol=0)
+
+        assert np.allclose(circ(params), -1)
+
+    @pytest.mark.jax
+    def test_jit(self):
+        """Test optimizer compatibility with jax.jit compilation."""
+        import jax
+        import jax.numpy as jnp
+
+        device = qml.device("default.qubit", wires=2)
+        qnode = qml.QNode(circuit, device=device)
+
+        params = [0.1, 0.2]
+        params = jnp.array(params)
+
+        opt = qml.QNGOptimizerQJIT()
+        state = opt.init(params)
+
+        new_params1, state1 = opt.step(qnode, params, state)
+        new_params2, state2, cost = opt.step_and_cost(qnode, params, state)
+
+        step = jax.jit(partial(opt.step, qnode))
+        step_and_cost = jax.jit(partial(opt.step_and_cost, qnode))
+        new_params1_jit, state1_jit = step(params, state)
+        new_params2_jit, state2_jit, cost_jit = step_and_cost(params, state)
+
+        # check params have been updated
+        assert not np.allclose(params, new_params1)
+        assert not np.allclose(params, new_params2)
+
+        # check jitted results match
+        assert np.allclose(new_params1, new_params1_jit)
+        assert state1 is None
+        assert state1_jit is None
+        assert np.allclose(new_params2, new_params2_jit)
+        assert state2 is None
+        assert state2_jit is None
+        assert np.allclose(cost, cost_jit)
+
+    @pytest.mark.jax
+    @pytest.mark.catalyst
+    @pytest.mark.external
+    def test_qjit(self):
+        """Test optimizer compatibility with qml.qjit compilation."""
+        import jax.numpy as jnp
+
+        pytest.importorskip("catalyst")
+
+        device = qml.device("lightning.qubit", wires=2)
+        qnode = qml.QNode(circuit, device=device)
+
+        params = [0.1, 0.2]
+        params = jnp.array(params)
+
+        opt = qml.QNGOptimizerQJIT()
+        state = opt.init(params)
+
+        new_params1, state1 = opt.step(qnode, params, state)
+        new_params2, state2, cost = opt.step_and_cost(qnode, params, state)
+
+        step = qml.qjit(partial(opt.step, qnode))
+        step_and_cost = qml.qjit(partial(opt.step_and_cost, qnode))
+        new_params1_qjit, state1_qjit = step(params, state)
+        new_params2_qjit, state2_qjit, cost_qjit = step_and_cost(params, state)
+
+        # check params have been updated
+        assert not np.allclose(params, new_params1)
+        assert not np.allclose(params, new_params2)
+
+        # check qjitted results match
+        assert np.allclose(new_params1, new_params1_qjit)
+        assert state1 is None
+        assert state1_qjit is None
+        assert np.allclose(new_params2, new_params2_qjit)
+        assert state2 is None
+        assert state2_qjit is None
+        assert np.allclose(cost, cost_qjit)