optax.contrib.MuonState

class optax.contrib.MuonState(

count: jax.typing.ArrayLike,

mu: optax.Updates,

ns_coeffs: jax.typing.ArrayLike,

)

State for the Muon algorithm.