optax.EmaState

class optax.EmaState(

count: jax.typing.ArrayLike,

ema: optax.Params,

)

Holds an exponential moving average of past updates.