optax.contrib.SophiaState# class optax.contrib.SophiaState(count: jax.Array, mu: base.Updates, nu: base.Updates, hessian_fn_state: Any)[source]# State for Sophia Optimizer.
