AbstractVariable#
- class penzai.core.variables.AbstractVariable[source]#
Bases:
ABCBase class for all variables.
Variables are allowed to be mutable, and should not be registered as PyTree nodes.
Methods
freeze()Returns a frozen copy of this variable.
get_slot()Returns the slot that this variable is replaced with when unbound.
update(new_frozen_value)Updates the value of this variable to match a frozen variable.
Inherited Methods
(expand to view inherited methods)
__init__()- final __get_state_as_jax_pytree__() AbstractVariableValue[source]#
Mutable variable protocol method, for JAX ecosystem interoperability.
- final __set_state_from_jax_pytree__(value: AbstractVariableValue)[source]#
Mutable variable protocol method, for JAX ecosystem interoperability.
- abstract freeze() AbstractVariableValue[source]#
Returns a frozen copy of this variable.
- abstract get_slot() AbstractVariableSlot[source]#
Returns the slot that this variable is replaced with when unbound.
The variable slot will uniquely identify this variable within a pytree.
- abstract update(new_frozen_value: AbstractVariableValue)[source]#
Updates the value of this variable to match a frozen variable.