jaxsnn.event.training.OptState

class jaxsnn.event.training.OptState(opt_state: optax.OptState, params: Dict[str, Optional[jax.Array]], rng: jax.Array)

Bases: NamedTuple

Container for the optimizer/training state across steps.

This immutable state groups the underlying optimizer’s internal state, the current set of model params, and the JAX PRNG key used for stochastic operations (e.g., sampling).

Parameters

• opt_state – Optimizer-specific internal state to carry across updates (e.g., from optax.init/optax.update).

• params – Ordered collection of learnable model parameters to be optimized.

• rng – JAX PRNG key used for randomized computations; should be split and updated between steps.

__init__()

Initialize self. See help(type(self)) for accurate signature.

Methods

Attributes

opt_state	Alias for field number 0
params	Alias for field number 1
rng	Alias for field number 2

opt_state: optax.OptState

Alias for field number 0

params: Dict[str, Optional[jax.Array]]

Alias for field number 1

rng: jax.Array

Alias for field number 2