attach_shared_parameters

penzai.nn.parameters.attach_shared_parameters(submodel: layer_base.Layer, *, strict: bool = True) → side_input.WithConstantSideInputs

Wraps a submodel in a handler that takes ownership of shareable params.

Penzai models must have only one copy of each parameter in the model pytree to ensure that gradients are computed properly. This function is responsible for transforming a model with multiple copies of a shareable parameter into a model with a single copy, using the SideInputEffect handler.

Parameters:

• submodel – Submodel containing ShareableUninitializedParameter instances that we should take ownership of.

• strict – If True, require that shared parameters have identical initializers by Python value. This should always be the case if the same uninitialized parameter instance was reused multiple times (e.g. if mark_shareable was used), and can catch issues if there are accidental name conflicts. If False, arbitrarily picks the first initializer seen for each shared parameter.

Returns:

A new pz.de.WithConstantSideInputs handler wraps the original submodel,

and holds a single UninitializedParameter for each uniquely-named ShareableUninitializedParameter in the original submodel.