optax.ScaleByAdamState

class optax.ScaleByAdamState(

count: jax.typing.ArrayLike,

mu: optax.Updates,

nu: optax.Updates,

)

State for the Adam algorithm.