def _build_selection_from_boundary(tree_with_boundary: Any) -> Selection:
  """Internal function to build a selection from a boundary.

  This function, like the `_InProgressSelectionBoundary` class, is an
  implementation detail and should not be relied upon by user code.

  Args:
    tree_with_boundary: A PyTree, with `_InProgressSelectionBoundary` nodes
      wrapping the subtrees that we wish to select. An unchecked requirement is
      that no `_InProgressSelectionBoundary` node should contain another
      `_InProgressSelectionBoundary` node; if this happens the boundary nodes
      will leak back into the returned selection.

  Returns:
    A selection, such that each occurrence of `_InProgressSelectionBoundary` is
    transformed into a selected subtree.
  """
  selected_by_path = collections.OrderedDict()

  def process_selected(keypath, leaf):
    """Process PyTree leaves or subtrees where _is_hole_or_quote is True."""
    if isinstance(leaf, _InProgressSelectionBoundary):
      # Select this object and stop processing the tree.
      selected_by_path[keypath] = leaf.selected
      return SelectionHole(keypath)
    elif _is_hole_or_quote(leaf):
      # This is already a hole or a quote, but we shouldn't treat it like a
      # hole or quote for the purposes of this selection. Escape them by adding
      # a level of quoting.
      return SelectionQuote(leaf)
    else:
      # This is just an ordinary leaf; keep it in the remainder.
      return leaf

  remainder = jax.tree_util.tree_map_with_path(
      process_selected, tree_with_boundary, is_leaf=_is_hole_or_quote
  )
  return Selection(selected_by_path=selected_by_path, remainder=remainder)
Internal function to build a selection from a boundary.

  This function, like the `_InProgressSelectionBoundary` class, is an
  implementation detail and should not be relied upon by user code.

  Args:
    tree_with_boundary: A PyTree, with `_InProgressSelectionBoundary` nodes
      wrapping the subtrees that we wish to select. An unchecked requirement is
      that no `_InProgressSelectionBoundary` node should contain another
      `_InProgressSelectionBoundary` node; if this happens the boundary nodes
      will leak back into the returned selection.

  Returns:
    A selection, such that each occurrence of `_InProgressSelectionBoundary` is
    transformed into a selected subtree.
  
penzai/core/selectors.py