def 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:
  """Makes a parameter variable (or slot) with a given name and initializer.

  Args:
    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.
  """
  if init_base_rng is None:
    return variables.ParameterSlot(name)
  else:
    if metadata is None:
      metadata = {}
    return variables.Parameter(
        value=initializer(
            derive_param_key(init_base_rng, name), *init_args, **init_kwargs
        ),
        label=name,
        metadata=metadata,
    )
Makes a parameter variable (or slot) with a given name and initializer.

  Args:
    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.
  
penzai/nn/parameters.py