def at_pytree_leaves(self) -> "Selection":
    """Selects all PyTree leaves of each selected subtree.

    This selects all of the leaves of the PyTree according to `jax.tree_util`,
    giving the most-specific selection expressible with a `Selection` object.
    (Note that, if any objects in the tree are not registered as JAX PyTree
    nodes, they will be selected in their entirety even if they contain children
    when printed out by treescope.)

    Returns:
      A new selection that selects every leaf of each selected subtree.
    """
    add_boundary = functools.partial(
        jax.tree_util.tree_map, _InProgressSelectionBoundary
    )
    return _build_selection_from_boundary(self.apply(add_boundary))
Selects all PyTree leaves of each selected subtree.

    This selects all of the leaves of the PyTree according to `jax.tree_util`,
    giving the most-specific selection expressible with a `Selection` object.
    (Note that, if any objects in the tree are not registered as JAX PyTree
    nodes, they will be selected in their entirety even if they contain children
    when printed out by treescope.)

    Returns:
      A new selection that selects every leaf of each selected subtree.
    
penzai/core/selectors.py