Implement new QNG Jax-based optimizer compatible with (q)jit

[SimoneGasperini]

For PL's non-quantum optimizers (e.g. GradientDescentOptimizer, AdamOptimizer, etc.), the user is currently redirected to leverage optax-native functionality (see here) 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:

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 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 Report

:white_check_mark: All modified and coverable lines are covered by tests. :white_check_mark: Project coverage is 99.68%. Comparing base (29e452d) to head (009fe13). :warning: Report is 446 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #7452   +/-   ##
=======================================
  Coverage   99.68%   99.68%           
=======================================
  Files         542      543    +1     
  Lines       55141    55204   +63     
=======================================
+ Hits        54967    55030   +63     
  Misses        174      174           

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

:rocket: New features to boost your workflow:

• :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[SimoneGasperini]

Thank you @PietropaoloFrisoni for taking a look!

are we waiting for the work to make qml.metric_tensor compatible with Catalyst to be completed before testing them?

qml.metric_tensor should be now compatible with Catalyst (with some limitations, see #7528) and its compatibility is tested here. However, it doesn't work yet with program capture enabled.

[josh146]

Nice! Don't forget to document the optimizer here: https://github.com/PennyLaneAI/pennylane/blob/efbb14f53f058c5a9a58f8c138e0e4bf0d81445e/doc/introduction/interfaces.rst?plain=1#L131

[SimoneGasperini]

Nice! Don't forget to document the optimizer here:

https://github.com/PennyLaneAI/pennylane/blob/efbb14f53f058c5a9a58f8c138e0e4bf0d81445e/doc/introduction/interfaces.rst?plain=1#L131

Thank you for the reminder @josh146! I added a few lines in 2f4ed35.