SharedParameterLookup

class penzai.nn.parameters.SharedParameterLookup

Bases: Struct, Generic[T]

A marker identifying a shared parameter.

A SharedParameterLookup acts like a parameter, but does not actually hold its value. Instead, it retrieves its parameter using a SideInputEffect, and expects the value of the parameter to be provided by some external handler, as configured by attach_shared_parameters.

Variables:

• ref (side_input.SideInputEffect[ParameterLike[T]]) – The SideInputEffect that provides the value for this parameter.

• value_structure (shapecheck.StructureAnnotation) – The structure of the value that will be substituted here.

Methods

__init__(ref, value_structure)
treescope_color()

Attributes

value	Value accessor for compatibility with ParameterLike.
ref
value_structure

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.

property value

Value accessor for compatibility with ParameterLike.

Raises:

MissingSharedParameterError – Since SharedParameterLookup does not have a value until given one by an attach_shared_parameters wrapper, retrieving their value always raises an error.