parameters

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:

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.