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)[source]#
State for schedule_free.