OptimState

class braintools.optim.OptimState(value, name=None, **metadata)

A brainstate.LongTermState holding optimizer state.

Used for quantities that persist across optimization steps but are not trainable parameters, such as optax optimizer state, step counters and learning rates.