def pretty_keystr(keypath: tuple[Any, ...], tree: Any) -> str:
  """Constructs a pretty name from a keypath and an object.

  This can be used to construct human-readable names for locations inside a
  PyTree.

  Args:
    keypath: Keypath to process.
    tree: Tree that this keypath indexes into.

  Returns:
    A human-readable string like
    "``Foo.sublayers[0]/Bar.body/Baz.param/Parameter.value"``
    instead of ``".sublayers[0].body.param.value"``
  """
  parts = []
  for key in keypath:
    if isinstance(
        key,
        jax.tree_util.GetAttrKey
        | jax.tree_util.FlattenedIndexKey
        | CustomGetAttrKey,
    ):
      parts.extend(("/", type(tree).__name__))
    split = tree_flatten_exactly_one_level(tree)
    assert split is not None
    tree = dict(split[0])[key]
    parts.append(str(key))
  if parts and parts[0] == "/":
    parts = parts[1:]
  return "".join(parts)
Constructs a pretty name from a keypath and an object.

  This can be used to construct human-readable names for locations inside a
  PyTree.

  Args:
    keypath: Keypath to process.
    tree: Tree that this keypath indexes into.

  Returns:
    A human-readable string like
    "``Foo.sublayers[0]/Bar.body/Baz.param/Parameter.value"``
    instead of ``".sublayers[0].body.param.value"``
  
penzai/core/tree_util.py