def random_split(
    key: jax.Array | NamedArrayBase,
    named_shape: Mapping[AxisName, int] | Sequence[tuple[AxisName, int]],
) -> NamedArray | NamedArrayView:
  """Splits a PRNG key into a `NamedArray` of PRNG keys with the given names.

  Args:
    key: PRNG key to split. Can also be a NamedArray of keys with disjoint names
      from ``named_shape``.
    named_shape: Named shape for the result. If an unordered mapping, the keys
      will be sorted before splitting. To avoid this (e.g. for unsortable keys)
      you can pass a `collections.OrderedDict` or a sequence of (name, size)
      tuples.

  Returns:
    `NamedArray` or `NamedArrayView` with the given named shape, filled with
    unique PRNG keys.
  """
  if not isinstance(key, NamedArrayBase):
    key = wrap(key)

  if isinstance(named_shape, collections.OrderedDict):
    names = named_shape.keys()
    sizes = named_shape.values()
  elif isinstance(named_shape, Sequence):
    names, sizes = zip(*named_shape)
  else:
    unsorted_keys = list(named_shape.keys())
    try:
      names = sorted(unsorted_keys)
    except Exception as exc:
      raise ValueError(
          "Unordered mappings must have sortable axis names when using"
          " `random_split`. If necessary, you can specify a particular ordering"
          " using a collections.OrderedDict or a tuple of (name, size) pairs."
      ) from exc
    sizes = [named_shape[name] for name in names]

  total_size = functools.reduce(operator.mul, sizes, 1)

  flat_split_keys = nmap(jax.random.split)(key, total_size)

  return flat_split_keys.reshape(
      tuple(sizes) + flat_split_keys.positional_shape[1:]
  ).tag_prefix(*names)
Splits a PRNG key into a `NamedArray` of PRNG keys with the given names.

  Args:
    key: PRNG key to split. Can also be a NamedArray of keys with disjoint names
      from ``named_shape``.
    named_shape: Named shape for the result. If an unordered mapping, the keys
      will be sorted before splitting. To avoid this (e.g. for unsortable keys)
      you can pass a `collections.OrderedDict` or a sequence of (name, size)
      tuples.

  Returns:
    `NamedArray` or `NamedArrayView` with the given named shape, filled with
    unique PRNG keys.
  
penzai/core/named_axes.py