def initialize_parameters_sharded(
    model: Any,
    prng_key: jax.Array,
    mesh: jax.sharding.Mesh,
    axis_name_to_mesh_name: dict[str, str | tuple[str, ...]] | None = None,
) -> Any:
  """Initializes the parameters of a model, sharded according to a mesh.

  Args:
    model: A model whose parameters we should initialize. Should usually contain
      `pz.nn.UninitializedParameter` instances.
    prng_key: Key to use to initialize parameters.
    mesh: The ``Mesh`` to shard the tree to.
    axis_name_to_mesh_name: A mapping from array axis names to mesh axis names.
      If an axis name is not present, that axis will not be sharded. If a mesh
      axis name is a tuple, the corresponding axis will be sharded to multiple
      mesh axes. If this dictionary is not provided, it will be inferred as an
      "identity" mapping, where each axis is sharded to a mesh axis with the
      same name (if present).

  Returns:
    An initialized version of `model`, with all uninitialized parameters
    replaced with initialized parameters, sharded according to the arguments.
  """
  uninit_params = (
      pz.select(model)
      .at_instances_of(pz.nn.UninitializedParameter)
      .get_sequence()
  )
  out_shapes = jax.eval_shape(
      pz.nn.initialize_parameters, uninit_params, prng_key
  )
  out_shardings = name_to_name_sharding(
      out_shapes, mesh, axis_name_to_mesh_name
  )
  init_fn = jax.jit(pz.nn.initialize_parameters, out_shardings=out_shardings)
  initialized_params = init_fn(uninit_params, prng_key)
  return (
      pz.select(model)
      .at_instances_of(pz.nn.UninitializedParameter)
      .set_sequence(initialized_params)
  )
Initializes the parameters of a model, sharded according to a mesh.

  Args:
    model: A model whose parameters we should initialize. Should usually contain
      `pz.nn.UninitializedParameter` instances.
    prng_key: Key to use to initialize parameters.
    mesh: The ``Mesh`` to shard the tree to.
    axis_name_to_mesh_name: A mapping from array axis names to mesh axis names.
      If an axis name is not present, that axis will not be sharded. If a mesh
      axis name is a tuple, the corresponding axis will be sharded to multiple
      mesh axes. If this dictionary is not provided, it will be inferred as an
      "identity" mapping, where each axis is sharded to a mesh axis with the
      same name (if present).

  Returns:
    An initialized version of `model`, with all uninitialized parameters
    replaced with initialized parameters, sharded according to the arguments.
  
penzai/deprecated/v1/toolshed/sharding_util.py