initialize_parameters

penzai.nn.parameters.initialize_parameters(model: Any, prng_key: jax.Array) → Any

Initializes all parameters in a model.

Parameters:

• model – Model to initialize.

• prng_key – Key to use to initialize parameters.

Returns:

Copy of model, where each instance of UninitializedParameter is replaced with a Parameter initialized using an independent PRNG key derived from prng_key. Keys are derived from the names of the parameters, so parameters with the same name will be initialized consistently if this method is given the same key.