Implement new MomentumQNGOptimizer
compatible with (q)jit and program capture
#7606
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The goal here is to implement a new version of the
MomentumQNGOptimizer
based on Jax and compatible withjax.jit
,qml.qjit
, and program capture. This PR depends on the q(jit)/capture compatible version of the standardQNGOptimizer
, 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 apennylane.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 nativeoptax
optimizers. This is required to avoid having a Jax tracer as a class attribute when using (q)jit compilation.[sc-92558]