optax.ScaleByBeliefState

class optax.ScaleByBeliefState(count: jax.typing.ArrayLike, mu: optax.Updates, nu: optax.Updates)

State for the rescaling by AdaBelief algorithm.