optax_adam

phasic.optax_wrapper.optax_adam(
    learning_rate=0.001,
    b1=0.9,
    b2=0.999,
    eps=1e-08,
)

Create Optax Adam optimizer wrapped for phasic.

Parameters

learning_rate : float or optax.Schedule = 0.001

Learning rate. Can be a scalar or an Optax schedule.

b1 : float = 0.9

Exponential decay rate for first moment.

b2 : float = 0.999

Exponential decay rate for second moment.

eps : float = 1e-8

Small constant for numerical stability.

Returns

: OptaxOptimizer

Wrapped Adam optimizer compatible with phasic SVGD.

Examples

>>> from phasic import optax_adam, SVGD
>>> svgd = SVGD(model=model, data=data, theta_dim=2, optimizer=optax_adam(0.001))