InternalTrainerState

class penzai.toolshed.basic_training.InternalTrainerState

Bases: Struct

Internal state for the trainer.

Variables:

• step (int | jax.Array) – Current step of training.

• opt_state (OptimizerStatePyTree) – An optimizer state for the learnable parts of model.

• loss_fn_state (LossStatePyTree) – Arbitrary state managed by the loss function. For instance, if your model has local state, you can functionalize it using pz.de.handle_local_states and store its state dict here.

Methods

__init__(step, opt_state, loss_fn_state)

Attributes

step
opt_state
loss_fn_state

Inherited Methods

(expand to view inherited methods)

attributes_dict()	Constructs a dictionary with all of the fields in the class.
from_attributes(**field_values)	Directly instantiates a struct given all of its fields.
key_for_field(field_name)	Generates a JAX PyTree key for a given field name.
select()	Wraps this struct in a selection, enabling functional-style mutations.
tree_flatten()	Flattens this tree node.
tree_flatten_with_keys()	Flattens this tree node with keys.
tree_unflatten(aux_data, children)	Unflattens this tree node.
treescope_color()	Computes a CSS color to display for this object in treescope.