zea.backend.optimizer

Simple implementation of optimizers that support multi-backend autodiff.

Functions

adam(step_size[, b1, b2, eps])

Construct optimizer triple for Adam.
zea.backend.optimizer.adam(step_size, b1=0.9, b2=0.999, eps=1e-08)[source]

Construct optimizer triple for Adam.

Implementation adapted from JAX’s example See example usage: JAX’s example usage

Parameters:

  • step_size – positive scalar

  • b1 – optional, a positive scalar value for beta_1, the exponential decay rate for the first moment estimates (default 0.9).

  • b2 – optional, a positive scalar value for beta_2, the exponential decay rate for the second moment estimates (default 0.999).

  • eps – optional, a positive scalar value for epsilon, a small constant for numerical stability (default 1e-8).

Returns:

An (init_fun, update_fun, get_params) triple.