AbstractVariableValue

class penzai.core.variables.AbstractVariableValue

Bases: Struct, ABC

Base class for all frozen variables.

Methods

get_slot()	Returns the slot that this variable is replaced with when unbound.
unfreeze_as_copy()	Returns a new mutable copy of this variable.

Inherited Methods

(expand to view inherited methods)

__init__()
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.

final __jax_pytree_state_to_new_variable__()

Mutable variable protocol method, for JAX ecosystem interoperability.

abstract get_slot() → AbstractVariableSlot

Returns the slot that this variable is replaced with when unbound.

The get_slot method of a frozen variable should always return the same slot as the get_slot method of the variable it was created from (and also for the variables created by unfreeze_as_copy).

abstract unfreeze_as_copy() → AbstractVariable

Returns a new mutable copy of this variable.