Skip to content

Implement new MomentumQNGOptimizer compatible with (q)jit and program capture #7606

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 11 commits into
base: master
Choose a base branch
from

Conversation

SimoneGasperini
Copy link
Contributor

@SimoneGasperini SimoneGasperini commented Jun 4, 2025

The goal here is to implement a new version of the MomentumQNGOptimizer based on Jax and compatible with jax.jit, qml.qjit, and program capture. This PR depends on the q(jit)/capture 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). 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]

@SimoneGasperini SimoneGasperini added the WIP 🚧 Work-in-progress label Jun 4, 2025
@SimoneGasperini SimoneGasperini marked this pull request as draft June 4, 2025 19:32
Copy link
Contributor

github-actions bot commented Jun 4, 2025

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
WIP 🚧 Work-in-progress
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant