quantax.optimizer.AdamUpdater

class quantax.optimizer.AdamUpdater

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.

If the new momentum values contain non-finite entries (e.g. from a non-finite gradient solve), the m, v and t buffers are left unchanged, so a single corrupted iteration can’t poison the accumulated momentum. The resulting non-finite step is still returned and is expected to be rejected by update.

__init__(mu: float = 0.95, beta: float = 0.995)

Parameters:

• mu – The first order momentum factor.

• beta – The second order momentum factor.

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

Initialize the persistent buffers for nparams parameters.

update(core_solve: Callable[[...], Array], Obar: Array, Ebar: Array, buffers: dict[str, Array]) → tuple[Array, dict[str, Array]]

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