optax_sgd

phasic.optax_wrapper.optax_sgd(learning_rate=0.01, momentum=0.0, nesterov=False)

Create Optax SGD optimizer wrapped for phasic.

Parameters

learning_rate : float or optax.Schedule = 0.01

Learning rate.

momentum : float = 0.0

Momentum coefficient. Set to 0.9 for standard momentum.

nesterov : bool = False

Whether to use Nesterov momentum.

Returns

: OptaxOptimizer

Wrapped SGD optimizer compatible with phasic SVGD.