optax.contrib.SAMState

class optax.contrib.SAMState(

steps_since_sync: jax.Array,

opt_state: base.OptState,

adv_state: base.OptState,

cache: base.Params | None,

)

State of GradientTransformation returned by sam.

steps_since_sync

Number of adversarial steps taken since the last outer update.

Type:

jax.Array

opt_state

State of the outer optimizer.

Type:

base.OptState

adv_state

State of the inner adversarial optimizer.

Type:

base.OptState

cache

a place to store the last outer updates.

Type:

Optional[base.Params]