make_parameter

Contents

make_parameter#

penzai.nn.parameters.make_parameter(name: str, init_base_rng: jax.Array | None, initializer: Callable[..., T], *init_args, metadata: dict[Any, Any] | None = None, **init_kwargs) → variables.Parameter[T] | variables.ParameterSlot[source]#

Makes a parameter variable (or slot) with a given name and initializer.

Parameters:

  • name – The name of the parameter.

  • init_base_rng – The base PRNG key to use for initialization, which will be combined with the name to obtain a unique key for this parameter. If None, the parameter will not be initialized; instead it will be set to a ParameterSlot to be filled in later.

  • initializer – The initializer function to use. The first argument should be a PRNG key, which will

  • *init_args – Positional arguments to pass to the initializer.

  • metadata – Metadata for the variable.

  • **init_kwargs – Keyword arguments to pass to the initializer.

Returns:

A Parameter if init_base_rng is not None, or a ParameterSlot otherwise.

Contents