quantax.optimizer.Adam#

class quantax.optimizer.Adam#

Bases: Updater

The AdamSR update, a variant of SR with first and second order momentum (like Adam) stored in the m, v and t buffers. The time cost is roughly twice of the plain update. The second order momentum enters as a diag_preconditioner argument: it is passed through to solvers that accept it (e.g. lsmr) and emulated by right preconditioning for those that don’t.

__init__(mu: float = 0.95, beta: float = 0.995, norm_clip: float | None = None)#
Parameters:

  • mu – The first order momentum factor.

  • beta – The second order momentum factor.

  • norm_clip – The maximum norm of the step to be accumulated. If not None, the raw step will be clipped to this value.

init(nparams: int) → dict[str, Array]#

Initialize the persistent buffers for nparams parameters.

update(core_solve, Obar, Ebar, buffers)#

Generate the optimization step and the new buffers, given the equation solver core_solve and the current buffers.