mark_shareable

penzai.nn.parameters.mark_shareable(submodel: Any)

Marks all uninitialized parameters in submodel as shareable.

This function is used to annotate the parameters in part of a model as being shareable: all UninitializedParameter instances will be replaced by equivalent ShareableUninitializedParameter instances. The submodel can then be used multiple times, as long as all uses are wrapped in a single call to attach_shared_parameters.

Parameter sharing is opt-in only to avoid accidentally sharing parameters due to name collisions. By convention, any caller that calls make_shareable is also responsible for constructing the attach_shared_parameters wrapper that will actually own the shared parameters.

Parameters:

submodel – Part of a model whose parameters should be shared.

Returns:

A copy of submodel where all UninitializedParameter instances have been replaced by ShareableUninitializedParameter instances.