optax.scale_by_stddev

optax.scale_by_stddev(decay: jax.typing.ArrayLike = 0.9, eps: jax.typing.ArrayLike = 1e-08, initial_scale: jax.typing.ArrayLike = 0.0, eps_in_sqrt: bool = True, bias_correction: bool = False) → optax.GradientTransformation

Rescale updates by the root of the centered exp. moving average of squares.

See optax.rmsprop() for more details.

Parameters:

• decay – Decay rate for the exponentially weighted average of squared grads.

• eps – Term added to the denominator to improve numerical stability.

• initial_scale – Initial value for second moment.

• eps_in_sqrt – Whether to add eps in the square root of the denominator or outside the square root.

• bias_correction – Whether to apply bias correction to the first and second moment.

Returns:

A optax.GradientTransformation object.