parameters

A simple parameter type.

Parameters in Penzai are identified by being instances of Parameter. By default, any array in a model that is not an instance of Parameter will be held constant during optimization.

Additionally, parameters have names, which can be used to identify a single parameter even after patching or reconfiguring a model. This is useful for storing parameters in checkpoints, for instance.

By convention, Penzai layers should annotate their parameters as being of type ParameterLike, which is a protocol that defines the interface a parameter should support. The main implementations of ParameterLike are:

• Parameter: An ordinary parameter that can be updated during optimization.

• FrozenParameter: A parameter that is frozen, which still conforms to the ParameterLike protocol, but will not be updated by default.

• UninitializedParameter: A parameter whose value has not yet been assigned. These are usually inserted when a model is being built the first time.

• ShareableUninitializedParameter: A temporary marker for an uninitialized parameter whose value will be shared by multiple layers. You usually don’t need to construct or manipulate these yourself; instead, use mark_shareable and attach_shared_parameters.

Additionally, users are free to substitute their own implementations of the ParameterLike protocol to customize how the value of a parameter should be constructed.

Neural network layers are responsible for assigning unique names to each parameter they create. The typical way to do this is to use add_parameter_prefix to add prefixes to the parameters inside any of their inner sublayers. Penzai does not automatically track name scopes, but it will raise errors if it detects multiple parameters with the same name.

Classes

FrozenParameter	A non-learnable parameter.
Parameter	A learnable parameter.
ParameterLike	Protocol for a parameter-like object.
ShareableUninitializedParameter	A shareable variant of an uninitialized parameter.
SharedParamTag	A tag for a shared parameter.
SharedParameterLookup	A marker identifying a shared parameter.
SupportsParameterRenaming	Base class that identifies a PyTree node as supporting parameter renaming.
UninitializedParameter	An uninitialized parameter.

Functions

add_parameter_prefix(prefix, tree[, delimiter])	Prepends a prefix to all parameter names inside the tree.
attach_shared_parameters(submodel, *[, strict])	Wraps a submodel in a handler that takes ownership of shareable params.
check_no_duplicated_parameters(model)	Checks that there are no duplicated parameters in a model.
initialize_parameters(model, prng_key)	Initializes all parameters in a model.
mark_shareable(submodel)	Marks all uninitialized parameters in submodel as shareable.

Exceptions

UninitializedParameterError

Raised when an uninitialized parameter is accessed.