optax.contrib.SophiaState

class optax.contrib.SophiaState(

count: jax.Array,

mu: base.Updates,

nu: base.Updates,

hessian_fn_state: Any,

)

State for Sophia Optimizer.