Class Selection

penzai/core/selectors.py:164–1411

A selected subset of nodes within a larger PyTree. Penzai selectors (such as ``.at(...)``) return ``Selection`` objects, which indicate a specific subset of nodes within a larger PyTree, allowing those nodes to be pulled out and modified in a functional way. Selected nodes are required to

@struct.pytree_dataclass
class Selection(Generic[SelectedSubtree], struct.Struct):
  """A selected subset of nodes within a larger PyTree.

  Penzai selectors (such as ``.at(...)``) return ``Selection`` objects, which
  indicate a specific subset of nodes within a larger PyTree, allowing those
  nodes to be pulled out and modified in a functional way.

  Selected nodes are required to be non-overlapping: no selected node can be
  the ancestor of any other selected node in the same selection.

  For convenience, a ``Selection`` is also a PyTree, and its leaves are the same
  as the leaves in the original PyTree, but they are likely to be in a different
  order.

  .. note:: Aside for functional programming geeks:
    A ``Selection`` is conceptually related to an "optic", specifically a
    "lens". If you're familiar with optics, you can
    think of a ``Selection`` as a partially-applied lens: it allows either
    retrieving the selected values, or setting the selected values in the
    structure. (If you're not familiar with optics, you can ignore this.)

  Attributes:
    selected_by_path: A mapping whose values are the selected parts from the
      original structure, and whose keys are the keypaths for those parts (as
      registered with JAX's PyTree registry). This is an `OrderedDict` to
      prevent JAX from trying to sort the keys, which may be arbitrarily
      hashable objects without an ordering.
    remainder: The rest of the structure. The locations where the selected
      components were are marked with `SelectionHole` nodes. If the remainder
      also includes a ``Selection`` itself, the remainder may also include
      `SelectionQuote` nodes.
  """

  selected_by_path: collections.OrderedDict[KeyPath, SelectedSubtree]
  remainder: Any

  def count(self) -> int:
    """Returns the number of elements in the selection."""
    return len(self.selected_by_path)

  def __len__(self) -> int:
    """Returns the number of elements in the selection."""
    return len(self.selected_by_path)

  def is_empty(self) -> bool:
    """Returns True if the selection is empty."""
    return not self.selected_by_path

  def deselect(self) -> Any:
    """Rebuilds the tree, forgetting which nodes were selected.

    Returns:
      A copy of `remainder` with the holes filled by the values in
      `selected_by_path`. If called on an ordinary selection, this rebuilds
      the original tree.
    """

    def rebuild(subtree) -> Any: