Repository: PennyLaneAI/pennylane
Pull request #7606 (2025-06-04 19:32:44)

Implement new `MomentumQNGOptimizer` compatible with (q)jit
The goal here is to implement a new Jax-based version of the `MomentumQNGOptimizer` compatible with `jax.jit` and `qml.qjit`. This PR depends on the q(jit) compatible version of the standard `QNGOptimizer` (see #7452 for more context).

Note that this optimization algorithm requires to keep track of a non-trivial internal state: in the current Autograd-based implementation, this state is stored in the optimizer's `accumulation` attribute as a `pennylane.numpy` array (see [here](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/optimize/momentum_qng.py#L122)). However, in this Jax-based version, the state needs to be carried over at each optimization step as it happens for the native `optax` optimizers. This is required to avoid having a Jax tracer as a class attribute when using (q)jit compilation.

[sc-92558]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7606?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `41.17647%` with `10 lines` in your changes missing coverage. Please review.
> Project coverage is 99.66%. Comparing base [(`8b6cacf`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/8b6cacfce96aadd46a03feba379d95c59189b458?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`6d45604`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/6d456042c500fbe69971b6b82cf0577893204361?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

| [Files with missing lines](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7606?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/optimize/momentum\_qng\_qjit.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7606?src=pr&el=tree&filepath=pennylane%2Foptimize%2Fmomentum_qng_qjit.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wdGltaXplL21vbWVudHVtX3FuZ19xaml0LnB5) | 37.50% | [10 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7606?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

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


```diff
@@            Coverage Diff             @@
##           master    #7606      +/-   ##
==========================================
- Coverage   99.68%   99.66%   -0.02%     
==========================================
  Files         543      544       +1     
  Lines       55205    55222      +17     
==========================================
+ Hits        55031    55038       +7     
- Misses        174      184      +10     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7606?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>
SimoneGasperini: @albi3ro @isaacdevlugt Code example updated!

Patch:
diff --git a/doc/introduction/interfaces.rst b/doc/introduction/interfaces.rst
index 257149bd5..fe92907f3 100644
--- a/doc/introduction/interfaces.rst
+++ b/doc/introduction/interfaces.rst
@@ -195,6 +195,7 @@ PennyLane offers a built-in Jax-based implementation compatible with ``jax.jit``
     :nosignatures:
 
     ~pennylane.QNGOptimizerQJIT
+    ~pennylane.MomentumQNGOptimizerQJIT
 
 :html:`</div>`
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 147c716bb..5b5d2ebe8 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,47 @@
 
 <h3>New features since last release</h3>
 
+* Leveraging quantum just-in-time compilation to optimize parameterized hybrid workflows with the momentum
+  quantum natural gradient optimizer is now possible with the new :class:`~.MomentumQNGOptimizerQJIT` optimizer.
+  [(#7606)](https://github.com/PennyLaneAI/pennylane/pull/7606)
+
+  Similar to the :class:`~.QNGOptimizerQJIT` optimizer, :class:`~.MomentumQNGOptimizerQJIT` offers a
+  `qml.qjit`-compatible analogue to the existing :class:`~.MomentumQNGOptimizer` with an Optax-like interface:
+
+  ```python
+  import pennylane as qml
+  import jax.numpy as jnp
+
+  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.MomentumQNGOptimizerQJIT(stepsize=0.1, momentum=0.2)
+
+  @qml.qjit
+  def update_step_qjit(i, args):
+      params, state = args
+      return opt.step(circuit, params, state)
+
+  @qml.qjit
+  def optimization_qjit(params, iters):
+      state = opt.init(params)
+      args = (params, state)
+      params, state = qml.for_loop(iters)(update_step_qjit)(args)
+      return params
+  ```
+
+  ```pycon
+  >>> params = jnp.array([0.1, 0.2])
+  >>> iters = 1000
+  >>> optimization_qjit(params=params, iters=iters)
+  Array([ 3.14159265, -1.57079633], dtype=float64)
+  ```
+
 <h3>Improvements 🛠</h3>
 
 <h4>OpenQASM-PennyLane interoperability</h4>
@@ -168,6 +209,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixes issue related to :func:`~.transforms.to_zx` adding the support for
+  `Toffoli` and `CCZ` gates conversion into their ZX-graph representation.
+  [(#7899)](https://github.com/PennyLaneAI/pennylane/pull/7899)
+
 * `get_best_diff_method` now correctly aligns with `execute` and `construct_batch` logic in workflows.
   [(#7898)](https://github.com/PennyLaneAI/pennylane/pull/7898)
 
@@ -190,6 +235,7 @@ Utkarsh Azad,
 Joey Carter,
 Yushao Chen,
 Marcus Edwards,
+Simone Gasperini,
 David Ittah,
 Erick Ochoa,
 Mudit Pandey,
diff --git a/pennylane/optimize/__init__.py b/pennylane/optimize/__init__.py
index 8ffe29ced..b549b3c5a 100644
--- a/pennylane/optimize/__init__.py
+++ b/pennylane/optimize/__init__.py
@@ -33,6 +33,7 @@ from .rotosolve import RotosolveOptimizer
 from .shot_adaptive import ShotAdaptiveOptimizer
 from .spsa import SPSAOptimizer
 from .qng_qjit import QNGOptimizerQJIT
+from .momentum_qng_qjit import MomentumQNGOptimizerQJIT
 
 # Optimizers to display in the docs
 __all__ = [
@@ -52,4 +53,5 @@ __all__ = [
     "ShotAdaptiveOptimizer",
     "SPSAOptimizer",
     "QNGOptimizerQJIT",
+    "MomentumQNGOptimizerQJIT",
 ]
diff --git a/pennylane/optimize/momentum_qng.py b/pennylane/optimize/momentum_qng.py
index caaf7b2ed..8744a3e3a 100644
--- a/pennylane/optimize/momentum_qng.py
+++ b/pennylane/optimize/momentum_qng.py
@@ -81,6 +81,8 @@ class MomentumQNGOptimizer(QNGOptimizer):
         - `QAOA <https://github.com/borbysh/Momentum-QNG/blob/main/QAOA_depth4.ipynb>`__
         - `VQE <https://github.com/borbysh/Momentum-QNG/blob/main/portfolio_optimization.ipynb>`__
 
+        See :class:`~.MomentumQNGOptimizerQJIT` for an Optax-like and ``jax.jit``/``qml.qjit``-compatible implementation.
+
     Keyword Args:
         stepsize=0.01 (float): the user-defined hyperparameter :math:`\eta`
         momentum=0.9 (float): the user-defined hyperparameter :math:`\rho`
diff --git a/pennylane/optimize/momentum_qng_qjit.py b/pennylane/optimize/momentum_qng_qjit.py
new file mode 100644
index 000000000..0280e244a
--- /dev/null
+++ b/pennylane/optimize/momentum_qng_qjit.py
@@ -0,0 +1,124 @@
+# 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 with momentum for Jax/Catalyst interface"""
+
+from pennylane import math
+
+from .qng_qjit import QNGOptimizerQJIT
+
+
+class MomentumQNGOptimizerQJIT(QNGOptimizerQJIT):
+    r"""Optax-like and ``jax.jit``/``qml.qjit``-compatible implementation of the :class:`~.MomentumQNGOptimizer`,
+    a generalized Quantum Natural Gradient (QNG) optimizer considering a discrete-time Langevin equation
+    with QNG force.
+
+    For more theoretical details, see the :class:`~.MomentumQNGOptimizer` documentation.
+
+    .. note::
+
+        Please be aware of the following:
+
+        - As with ``MomentumQNGOptimizer``, ``MomentumQNGOptimizerQJIT`` supports a single QNode to encode the objective function.
+
+        - ``MomentumQNGOptimizerQJIT`` does not support any QNode with multiple arguments. A potential workaround
+          would be to combine all parameters into a single objective function argument.
+
+        - ``MomentumQNGOptimizerQJIT`` does not work correctly if there is any classical processing in the QNode circuit
+          (e.g., ``2 * theta`` as a gate parameter).
+
+    Parameters:
+        stepsize (float): the stepsize hyperparameter (default value: 0.01).
+        momentum (float): the momentum coefficient hyperparameter (default value: 0.9).
+        approx (str): approximation method for the metric tensor (default value: "block-diag").
+
+            - 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 (float): metric tensor regularization to be applied at each optimization step (default value: 0).
+
+    **Example:**
+
+    Consider a hybrid workflow to optimize an objective function defined by a quantum circuit.
+    To make the entire workflow faster, the update step and the whole optimization
+    can be just-in-time compiled using the :func:`~.qjit` decorator:
+
+    .. code-block:: python
+
+        import pennylane as qml
+        import jax.numpy as jnp
+
+        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.MomentumQNGOptimizerQJIT(stepsize=0.1, momentum=0.2)
+
+        @qml.qjit
+        def update_step_qjit(i, args):
+            params, state = args
+            return opt.step(circuit, params, state)
+
+        @qml.qjit
+        def optimization_qjit(params, iters):
+            state = opt.init(params)
+            args = (params, state)
+            params, state = qml.for_loop(iters)(update_step_qjit)(args)
+            return params
+
+    >>> params = jnp.array([0.1, 0.2])
+    >>> iters = 1000
+    >>> optimization_qjit(params=params, iters=iters)
+    Array([ 3.14159265, -1.57079633], dtype=float64)
+
+    Make sure you are using the ``lightning.qubit`` device along with ``qml.qjit``.
+    """
+
+    def __init__(self, stepsize=0.01, momentum=0.9, approx="block-diag", lam=0):
+        super().__init__(stepsize, approx, lam)
+        self.momentum = momentum
+
+    def init(self, params):
+        """Return the initial state of the optimizer. This state is always initialized as an
+        array of zeros with the same shape and type of the given array of parameters.
+
+        Args:
+            params (array): QNode parameters
+
+        Returns:
+            array: initial state of the optimizer
+        """
+        # pylint:disable=no-self-use
+        return math.zeros_like(params)
+
+    def _apply_grad(self, mt, grad, params, state):
+        """Update the optimizer's state and the array of parameters for a single optimization
+        step according to the Quantum Natural Gradient algorithm with momentum.
+        """
+        shape = math.shape(grad)
+        grad_flat = math.flatten(grad)
+        update_flat = math.linalg.pinv(mt) @ grad_flat
+        update = math.reshape(update_flat, shape)
+        state = self.momentum * state + self.stepsize * update
+        new_params = params - state
+        return new_params, state
diff --git a/pennylane/optimize/qng_qjit.py b/pennylane/optimize/qng_qjit.py
index 7bd04744a..f49109fba 100644
--- a/pennylane/optimize/qng_qjit.py
+++ b/pennylane/optimize/qng_qjit.py
@@ -78,12 +78,13 @@ class QNGOptimizerQJIT:
     Array([ 3.14159265, -1.57079633], dtype=float64)
 
     Make sure you are using the ``lightning.qubit`` device along with ``qml.qjit`` with ``autograph`` enabled.
+    Using ``qml.qjit`` on the whole workflow with ``autograph`` not enabled may lead to a substantial increase
+    in compilation time and no runtime benefits.
 
-    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:
+    The ``jax.jit`` decorator should not be used on the entire workflow.
+    However, it can be used with the ``default.qubit`` device to just-in-time
+    compile the ``step`` (or ``step_and_cost``) method of the optimizer, leading
+    to a significative increase in runtime performance:
 
     .. code-block:: python
 
@@ -163,13 +164,6 @@ class QNGOptimizerQJIT:
 
         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)
@@ -188,13 +182,6 @@ class QNGOptimizerQJIT:
 
         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)
diff --git a/pennylane/transforms/zx/converter.py b/pennylane/transforms/zx/converter.py
index 688065962..5a2d2b397 100644
--- a/pennylane/transforms/zx/converter.py
+++ b/pennylane/transforms/zx/converter.py
@@ -29,6 +29,52 @@ from pennylane.typing import PostprocessingFn
 from pennylane.wires import Wires
 
 
+def _toffoli_clifford_t_decomp(wires):
+    """Return the explicit Clifford+T decomposition of the Toffoli gate,
+    replacing Adjoint(T) with PhaseShift(-π/4)."""
+
+    return [
+        qml.Hadamard(wires=wires[2]),
+        qml.CNOT(wires=[wires[1], wires[2]]),
+        qml.PhaseShift(-np.pi / 4, wires=wires[2]),
+        qml.CNOT(wires=[wires[0], wires[2]]),
+        qml.T(wires=wires[2]),
+        qml.CNOT(wires=[wires[1], wires[2]]),
+        qml.PhaseShift(-np.pi / 4, wires=wires[2]),
+        qml.CNOT(wires=[wires[0], wires[2]]),
+        qml.T(wires=wires[2]),
+        qml.T(wires=wires[1]),
+        qml.CNOT(wires=[wires[0], wires[1]]),
+        qml.Hadamard(wires=wires[2]),
+        qml.T(wires=wires[0]),
+        qml.PhaseShift(-np.pi / 4, wires=wires[1]),
+        qml.CNOT(wires=[wires[0], wires[1]]),
+    ]
+
+
+def _ccz_clifford_t_decomp(wires):
+    """Return the explicit Clifford+T decomposition of the CCZ gate,
+    replacing Adjoint(T) with PhaseShift(-π/4)."""
+
+    return [
+        qml.CNOT(wires=[wires[1], wires[2]]),
+        qml.PhaseShift(-np.pi / 4, wires=wires[2]),
+        qml.CNOT(wires=[wires[0], wires[2]]),
+        qml.T(wires=wires[2]),
+        qml.CNOT(wires=[wires[1], wires[2]]),
+        qml.PhaseShift(-np.pi / 4, wires=wires[2]),
+        qml.CNOT(wires=[wires[0], wires[2]]),
+        qml.T(wires=wires[2]),
+        qml.T(wires=wires[1]),
+        qml.CNOT(wires=[wires[0], wires[1]]),
+        qml.Hadamard(wires=wires[2]),
+        qml.T(wires=wires[0]),
+        qml.PhaseShift(-np.pi / 4, wires=wires[1]),
+        qml.CNOT(wires=[wires[0], wires[1]]),
+        qml.Hadamard(wires=wires[2]),
+    ]
+
+
 class VertexType:  # pylint: disable=too-few-public-methods
     """Type of a vertex in the graph.
 
@@ -254,6 +300,13 @@ def to_zx(tape, expand_measurements=False):
         It is a PennyLane adapted and reworked `circuit_to_graph <https://github.com/Quantomatic/pyzx/blob/master/pyzx/circuit/graphparser.py>`_
         function.
 
+    .. note::
+
+        Prior to being added to the graph, Toffoli and CCZ gates are replaced by particular decompositions. These decompositions
+        are described in detail in: J. Welch, A. Bocharov, and K. Svore, “Efficient Approximation of Diagonal Unitaries over the Clifford+T Basis,”
+        Quantum information & computation, vol. 16, Dec. 2014, doi: 10.26421/QIC16.1-2-6.
+        This is necessary because Toffoli and CCZ gates are not directly supported in PyZX.
+
         Copyright (C) 2018 - Aleks Kissinger and John van de Wetering
     """
     # If it is a simple operation just transform it to a tape
@@ -284,17 +337,22 @@ def _to_zx_transform(
 
     # Dictionary of gates (PennyLane to PyZX circuit)
     gate_types = {
-        "PauliX": pyzx.circuit.gates.NOT,
-        "PauliZ": pyzx.circuit.gates.Z,
+        "X": pyzx.circuit.gates.NOT,
+        "Y": pyzx.circuit.gates.Y,
+        "Z": pyzx.circuit.gates.Z,
         "S": pyzx.circuit.gates.S,
         "T": pyzx.circuit.gates.T,
         "Hadamard": pyzx.circuit.gates.HAD,
         "RX": pyzx.circuit.gates.XPhase,
+        "RY": pyzx.circuit.gates.YPhase,
         "RZ": pyzx.circuit.gates.ZPhase,
         "PhaseShift": pyzx.circuit.gates.ZPhase,
         "SWAP": pyzx.circuit.gates.SWAP,
         "CNOT": pyzx.circuit.gates.CNOT,
+        "CY": pyzx.circuit.gates.CY,
         "CZ": pyzx.circuit.gates.CZ,
+        "CRX": pyzx.circuit.gates.CRX,
+        "CRY": pyzx.circuit.gates.CRY,
         "CRZ": pyzx.circuit.gates.CRZ,
         "CH": pyzx.circuit.gates.CHAD,
         "CCZ": pyzx.circuit.gates.CCZ,
@@ -331,17 +389,12 @@ def _to_zx_transform(
         )
 
         expanded_operations = []
-
-        # Define specific decompositions
         for op in mapped_tape.operations:
-            if op.name == "RY":
-                theta = op.data[0]
-                decomp = [
-                    qml.RX(np.pi / 2, wires=op.wires),
-                    qml.RZ(theta + np.pi, wires=op.wires),
-                    qml.RX(np.pi / 2, wires=op.wires),
-                    qml.RZ(3 * np.pi, wires=op.wires),
-                ]
+            if op.name == "Toffoli":
+                decomp = _toffoli_clifford_t_decomp(op.wires)
+                expanded_operations.extend(decomp)
+            elif op.name == "CCZ":
+                decomp = _ccz_clifford_t_decomp(op.wires)
                 expanded_operations.extend(decomp)
             else:
                 expanded_operations.append(op)


Test patch:
diff --git a/tests/optimize/test_momentum_qng_qjit.py b/tests/optimize/test_momentum_qng_qjit.py
new file mode 100644
index 000000000..0643de8af
--- /dev/null
+++ b/tests/optimize/test_momentum_qng_qjit.py
@@ -0,0 +1,240 @@
+# 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 Momentum-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 MomentumQNGOptimizerQJIT."""
+
+    def test_initialization_default(self):
+        """Test that initializing MomentumQNGOptimizerQJIT with default values works."""
+        opt = qml.MomentumQNGOptimizerQJIT()
+        assert opt.stepsize == 0.01
+        assert opt.momentum == 0.9
+        assert opt.approx == "block-diag"
+        assert opt.lam == 0
+
+    def test_initialization_custom(self):
+        """Test that initializing MomentumQNGOptimizerQJIT with custom values works."""
+        opt = qml.MomentumQNGOptimizerQJIT(stepsize=0.05, momentum=0.8, approx="diag", lam=1e-9)
+        assert opt.stepsize == 0.05
+        assert opt.momentum == 0.8
+        assert opt.approx == "diag"
+        assert opt.lam == 1e-9
+
+    def test_init_zero_state(self):
+        """Test that the MomentumQNGOptimizerQJIT state is initialized to an array of zeros."""
+        opt = qml.MomentumQNGOptimizerQJIT()
+        params = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
+        state = opt.init(params)
+        assert np.all(state == np.zeros_like(params))
+
+
+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))
+
+        opt = qml.MomentumQNGOptimizerQJIT(stepsize=0.1, momentum=0.4)
+        params = jnp.array([0.11, 0.412])
+        state = jnp.array([[-0.31, 0.842]])
+
+        new_params1, state1 = opt.step(circ, params, state)
+        new_params2, state2, cost = opt.step_and_cost(circ, params, state)
+
+        exp_params = qml.numpy.array(params)
+        exp_state = np.array([-0.31, 0.842])
+        exp_cost = circ(exp_params)
+        exp_mt = np.array([0.25, (np.cos(exp_params[0]) ** 2) / 4])
+        exp_state = opt.momentum * exp_state + opt.stepsize * qml.grad(circ)(exp_params) / exp_mt
+        exp_params -= exp_state
+
+        assert np.allclose(new_params1, exp_params)
+        assert np.allclose(new_params2, exp_params)
+        assert np.allclose(state1, exp_state)
+        assert np.allclose(state2, exp_state)
+        assert np.allclose(cost, exp_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, wires=4))
+        def circ(params):
+            qml.DoubleExcitation(params[0], wires=[0, 1, 2, 3])
+            qml.RY(params[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        opt = qml.MomentumQNGOptimizerQJIT(stepsize=0.1, momentum=0.4)
+        params = jnp.array([0.11, 0.412])
+        state = jnp.array([[-0.31, 0.842]])
+
+        new_params1, state1 = opt.step(circ, params, state)
+        new_params2, state2, cost = opt.step_and_cost(circ, params, state)
+
+        exp_params = qml.numpy.array(params)
+        exp_state = np.array([-0.31, 0.842])
+        exp_cost = circ(exp_params)
+        exp_mt = np.array([1 / 16, 1 / 4])
+        exp_state = opt.momentum * exp_state + opt.stepsize * qml.grad(circ)(exp_params) / exp_mt
+        exp_params -= exp_state
+
+        assert np.allclose(new_params1, exp_params)
+        assert np.allclose(new_params2, exp_params)
+        assert np.allclose(state1, exp_state)
+        assert np.allclose(state2, exp_state)
+        assert np.allclose(cost, exp_cost)
+
+    @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.MomentumQNGOptimizerQJIT(stepsize=0.2, momentum=0.3)
+        params = jnp.array([0.011, 0.012])
+        state = opt.init(params)
+
+        exp_params = np.array([0.011, 0.012])
+        exp_state = np.zeros_like(exp_params)
+
+        num_steps = 30
+        for _ in range(num_steps):
+            params, state, cost = opt.step_and_cost(circ, params, state)
+
+            exp_cost = circ(exp_params)
+            exp_mt = np.array([0.25, (np.cos(exp_params[0]) ** 2) / 4])
+            exp_state = opt.momentum * exp_state + opt.stepsize * grad(exp_params) / exp_mt
+            exp_params -= exp_state
+
+            assert np.allclose(cost, exp_cost, atol=tol, rtol=0)
+            assert np.allclose(state, exp_state, atol=tol, rtol=0)
+            assert np.allclose(params, exp_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)
+
+        opt = qml.MomentumQNGOptimizerQJIT()
+        params = jnp.array([0.1, 0.2])
+        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(state, state1)
+        assert not np.allclose(params, new_params2)
+        assert not np.allclose(state, state2)
+
+        # check jitted results match
+        assert np.allclose(new_params1, new_params1_jit)
+        assert np.allclose(state1, state1_jit)
+        assert np.allclose(new_params2, new_params2_jit)
+        assert np.allclose(state2, state2_jit)
+        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)
+
+        opt = qml.MomentumQNGOptimizerQJIT()
+        params = jnp.array([0.1, 0.2])
+        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 and state have been updated
+        assert not np.allclose(params, new_params1)
+        assert not np.allclose(state, state1)
+        assert not np.allclose(params, new_params2)
+        assert not np.allclose(state, state2)
+
+        # check qjitted results match
+        assert np.allclose(new_params1, new_params1_qjit)
+        assert np.allclose(state1, state1_qjit)
+        assert np.allclose(new_params2, new_params2_qjit)
+        assert np.allclose(state2, state2_qjit)
+        assert np.allclose(cost, cost_qjit)
diff --git a/tests/transforms/test_zx.py b/tests/transforms/test_zx.py
index 733cbc0f9..c8a6f82db 100644
--- a/tests/transforms/test_zx.py
+++ b/tests/transforms/test_zx.py
@@ -26,31 +26,30 @@ from pennylane.tape import QuantumScript
 from pennylane.transforms import TransformError
 
 pyzx = pytest.importorskip("pyzx")
-
 pytestmark = pytest.mark.external
-
 supported_operations = [
-    qml.PauliX(wires=0),
-    qml.PauliZ(wires=0),
-    qml.Hadamard(wires=0),
+    qml.X(wires=0),
+    qml.Y(wires=0),
+    qml.Z(wires=0),
     qml.S(wires=0),
     qml.T(wires=0),
+    qml.Hadamard(wires=0),
     qml.SWAP(wires=[0, 1]),
     qml.CNOT(wires=[0, 1]),
+    qml.CY(wires=[0, 1]),
     qml.CZ(wires=[0, 1]),
     qml.CH(wires=[0, 1]),
-]
-
-supported_operations_params = [
     qml.RX(0.3, wires=0),
+    qml.RY(0.3, wires=0),
     qml.RZ(0.3, wires=0),
+    qml.PhaseShift(0.3, wires=0),
+    qml.CRX(0.3, wires=[0, 1]),
+    qml.CRY(0.3, wires=[0, 1]),
     qml.CRZ(0.3, wires=[0, 1]),
+    qml.Toffoli(wires=[0, 1, 2]),
+    qml.CCZ(wires=[0, 1, 2]),
 ]
 
-expanded_operations = [qml.PauliY(wires=0), qml.PhaseShift(0.3, wires=0), qml.RY(0.3, wires=0)]
-
-non_diagram_like_operations = [qml.CCZ(wires=[0, 1, 2]), qml.Toffoli(wires=[0, 1, 2])]
-
 decompose_phases = [True, False]
 qscript = [True, False]
 
@@ -81,7 +80,7 @@ class TestConvertersZX:
 
     @pytest.mark.parametrize("script", qscript)
     @pytest.mark.parametrize("operation", supported_operations)
-    def test_supported_operation_no_params(self, operation, script):
+    def test_supported_operations(self, operation, script):
         """Test to convert the script to a ZX graph and back for supported operations."""
 
         I = qml.math.eye(2 ** len(operation.wires))
@@ -118,109 +117,11 @@ class TestConvertersZX:
 
         assert qml.math.allclose(mat_product, I)
 
-    @pytest.mark.parametrize("script", qscript)
-    @pytest.mark.parametrize("operation", supported_operations_params)
-    def test_supported_operation_params(self, operation, script):
-        """Test to convert the script to a ZX graph and back for supported operations with parameters."""
-
-        if script:
-            qs = QuantumScript([operation])
-        else:
-            qs = operation
-
-        I = qml.math.eye(2 ** len(operation.wires))
-
-        matrix_qscript = qml.matrix(qs, wire_order=qs.wires)
-
-        zx_g = qml.transforms.to_zx(qs)
-
-        matrix_zx = zx_g.to_matrix()
-
-        assert isinstance(zx_g, pyzx.graph.graph_s.GraphS)
-        # Check whether the two matrices are each others conjugate transposes
-        mat_product = qml.math.dot(matrix_qscript, qml.math.conj(matrix_zx.T))
-        # Remove global phase
-        mat_product /= mat_product[0, 0]
-
-        assert qml.math.allclose(mat_product, I)
-
-        qscript_back = qml.transforms.from_zx(zx_g)
-        assert isinstance(qscript_back, qml.tape.QuantumScript)
-
-        matrix_qscript_back = qml.matrix(qscript_back, wire_order=list(range(len(qs.wires))))
-
-        # Check whether the two matrices are each others conjugate transposes
-        mat_product = qml.math.dot(matrix_qscript, qml.math.conj(matrix_qscript_back.T))
-        # Remove global phase
-        mat_product /= mat_product[0, 0]
-
-        assert qml.math.allclose(mat_product, I)
-
-    @pytest.mark.parametrize("script", qscript)
-    @pytest.mark.parametrize("operation", expanded_operations)
-    def test_operation_need_expansion(self, operation, script):
-        """Test to convert the script to a ZX graph and back for operations that needs expansions."""
-
-        if script:
-            qs = QuantumScript([operation])
-        else:
-            qs = operation
-
-        I = qml.math.eye(2 ** len(operation.wires))
-
-        matrix_qscript = qml.matrix(qs, wire_order=qs.wires)
-
-        zx_g = qml.transforms.to_zx(qs)
-
-        matrix_zx = zx_g.to_matrix()
-
-        assert isinstance(zx_g, pyzx.graph.graph_s.GraphS)
-        # Check whether the two matrices are each others conjugate transposes
-        mat_product = qml.math.dot(matrix_qscript, qml.math.conj(matrix_zx.T))
-        # Remove global phase
-        mat_product /= mat_product[0, 0]
-
-        assert qml.math.allclose(mat_product, I)
-
-        qscript_back = qml.transforms.from_zx(zx_g)
-        assert isinstance(qscript_back, qml.tape.QuantumScript)
-
-        matrix_qscript_back = qml.matrix(qscript_back, wire_order=list(range(len(qs.wires))))
-
-        # Check whether the two matrices are each others conjugate transposes
-        mat_product = qml.math.dot(matrix_qscript, qml.math.conj(matrix_qscript_back.T))
-        # Remove global phase
-        mat_product /= mat_product[0, 0]
-
-        assert qml.math.allclose(mat_product, I)
-
-    @pytest.mark.parametrize("operation", non_diagram_like_operations)
-    def test_non_diagram_like_op(self, operation):
-        """Test operations that result in a non diagram like circuit."""
-
-        I = qml.math.eye(2 ** len(operation.wires))
-
-        qs = QuantumScript([operation], [])
-        matrix_qscript = qml.matrix(qs, wire_order=qs.wires)
-
-        zx_g = qml.transforms.to_zx(qs)
-        assert isinstance(zx_g, pyzx.graph.graph_s.GraphS)
-
-        matrix_zx = zx_g.to_matrix()
-        # Check whether the two matrices are each others conjugate transposes
-        mat_product = qml.math.dot(matrix_qscript, qml.math.conj(matrix_zx.T))
-        # Remove global phase
-        mat_product /= mat_product[0, 0]
-        assert qml.math.allclose(mat_product, I)
-
-        with pytest.raises(QuantumFunctionError, match="Graph doesn't seem circuit like"):
-            qml.transforms.from_zx(zx_g)
-
     @pytest.mark.parametrize("decompose", decompose_phases)
     def test_circuit(self, decompose):
         """Test a simple circuit."""
 
-        I = qml.math.eye(2**2)
+        I = qml.math.eye(2**3)
 
         operations = [
             qml.RZ(5 / 4 * np.pi, wires=0),
@@ -234,6 +135,8 @@ class TestConvertersZX:
             qml.CNOT(wires=[0, 1]),
             qml.CNOT(wires=[1, 0]),
             qml.SWAP(wires=[0, 1]),
+            qml.Toffoli(wires=[0, 1, 2]),
+            qml.CCZ(wires=[0, 1, 2]),
         ]
 
         qs = QuantumScript(operations, [])