StateVariable

StateVariable#

class penzai.core.variables.StateVariable[source]#

Bases: LabeledVariable[T]

A mutable state variable.

StateVariables are variables that can be updated inside the model (and are usually not updated by optimizers).

Variables:

  • label (VariableLabel) – The unique label for this variable.

  • value – The mutable value stored in the variable. Should be a JAX pytree.

  • metadata (dict[Any, Any]) – A dictionary of metadata associated with this parameter.

Inherited Attributes

value

Methods

__init__(value, *[, label, metadata])

Constructs a new state variable.

freeze()

get_slot()

treescope_color()

Attributes

value

label

metadata

Inherited Methods

(expand to view inherited methods)

set_value(new_val)

Sets the value of the Variable.

update(new_frozen_value)

Updates the value of this variable to match a frozen variable.
__init__(value: T, *, label: VariableLabel | None = None, metadata: dict[Any, Any] | None = None)[source]#

Constructs a new state variable.

Parameters:

  • value – The initial value of the variable.

  • label – The unique label for this variable. If not provided, a unique label will be generated automatically.

  • metadata – A dictionary of metadata associated with this Variable.

Contents