def derive_param_key(base_rng: jax.Array, name: str) -> jax.Array:
  """Derives a PRNG key for a parameter from a base key and a name.

  Args:
    base_rng: The base PRNG key.
    name: The name of the parameter.

  Returns:
    A unique PRNG key for the parameter with this name.
  """
  hashed_bytes = hashlib.sha1(name.encode("utf-8")).digest()
  offset = int.from_bytes(hashed_bytes[:4], byteorder="big")
  return jax.random.fold_in(base_rng, offset)
Derives a PRNG key for a parameter from a base key and a name.

  Args:
    base_rng: The base PRNG key.
    name: The name of the parameter.

  Returns:
    A unique PRNG key for the parameter with this name.
  
penzai/nn/parameters.py