optax_rmsprop

phasic.optax_wrapper.optax_rmsprop(
    learning_rate=0.001,
    decay=0.9,
    eps=1e-08,
    momentum=0.0,
)

Create Optax RMSprop optimizer wrapped for phasic.

Parameters

learning_rate : float or optax.Schedule = 0.001

Learning rate.

decay : float = 0.9

Decay rate for exponential moving average.

eps : float = 1e-8

Small constant for numerical stability.

momentum : float = 0.0

Momentum coefficient.

Returns

: OptaxOptimizer

Wrapped RMSprop optimizer compatible with phasic SVGD.