optax.contrib.ScheduleFreeState

class optax.contrib.ScheduleFreeState(

b1: jax.typing.ArrayLike,

weight_sum: jax.typing.ArrayLike,

step_count: jax.typing.ArrayLike,

max_lr: jax.typing.ArrayLike,

base_optimizer_state: optax.OptState,

z: optax.Params,

)

State for schedule_free.