optax.LookaheadState

class optax.LookaheadState(

fast_state: optax.OptState,

steps_since_sync: Array,

)

State of the GradientTransformation returned by lookahead.

fast_state

Optimizer state of the fast optimizer.

Type:

base.OptState

steps_since_sync

Number of fast optimizer steps taken since slow and fast parameters were synchronized.

Type:

jax.Array