parameters

Parameters and state variables in Penzai neural networks.

Parameters in Penzai models are allowed to be any object that has a value attribute, as described in the ParameterLike protocol. This makes it possible to wrap parameters or apply preprocessing logic when values are retrieved. Usually, however, Penzai parameters are instances of Parameter, which is a subclass of Variable. This allows them to be easily shared and updated in place.

By convention, neural network initializers take an argument init_base_rng, which is a JAX PRNG key that is shared across the full model. They then initialize their parameters by combining the init_base_rng with the unique string name of each parameter.

Classes

ParameterLike

Protocol for a parameter-like object.

Functions

assert_no_parameter_slots(model)	Asserts that the given model has no ParameterSlot subtrees.
derive_param_key(base_rng, name)	Derives a PRNG key for a parameter from a base key and a name.
make_parameter(name, init_base_rng, ...[, ...])	Makes a parameter variable (or slot) with a given name and initializer.