def initialize_parameters(model: Any, prng_key: jax.Array) -> Any:
  """Initializes all parameters in a model.

  Args:
    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.
  """
  check_no_duplicated_parameters(model)

  def name_to_key(name: str):
    hashed_bytes = hashlib.sha1(name.encode("utf-8")).digest()
    offset = int.from_bytes(hashed_bytes[:4], byteorder="big")
    return jax.random.fold_in(prng_key, offset)

  return (
      selectors.select(model)
      .at_instances_of(UninitializedParameter)
      .apply(lambda uninit: uninit.initialize(name_to_key(uninit.name)))
  )
Initializes all parameters in a model.

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