Method at

penzai/core/selectors.py:501–629

Selects a specific child of each selected node. ``Selection.at(...)`` allows you to modify a tree with an almost-imperative style while maintaining a functional interface, similar to the ``Array.at[...]`` syntax for ordinary NDArrays. It takes a callable that picks out a subtree of

(
      self,
      accessor_fn: Callable[[SelectedSubtree], Any | Collection[Any]],
      multiple: bool | None = None,
  )

    ).flatten_selected_selections()

  def at(
      self,
      accessor_fn: Callable[[SelectedSubtree], Any | Collection[Any]],
      multiple: bool | None = None,
  ) -> "Selection":
    """Selects a specific child of each selected node.

    ``Selection.at(...)`` allows you to modify a tree with an almost-imperative
    style while maintaining a functional interface, similar to the
    ``Array.at[...]`` syntax for ordinary NDArrays. It takes a callable that
    picks out a subtree of the tree, and returns a new selection that selects
    the part that was picked out.

    For instance, if you have an object

    ::
      obj = Foo(bar=[1, 2, {"baz": 5}])

    you could select the 5 using

    ::

      pz.select(obj).at(lambda x: x.bar[2]["baz"])

    Args:
      accessor_fn: A function which takes each element of the current selection
        and returns a single node within that selection (if ``multiple`` is
        False) or a collection of nodes (if ``multiple`` is True). This function
        must be structural; it must depend only on the PyTree structure of its
        input and not on the actual values or Python IDs of the leaves. It will
        be called with a copy of the object where every PyTree leaf and every
        empty PyTree node (e.g. an empty tuple or the None singleton) are
        wrapped with an internal wrapper object.
      multiple: Whether `accessor_fn` returns a collection of nodes to select,
        rather than a single node. If `None`, first tries to find it as a single
        node, and if that fails, tries to find it as a collection of nodes but
        emits a warning.

    Returns:
      A modified selection that selects the specific child of each node in the
      original selection (or the set of nodes if ``multiple`` was True).
    """

    def _is_leaf_or_childless(node):
      result = penzai_tree_util.tree_flatten_exactly_one_level(node)
      if result is None:
        return True
      else:
        children, _ = result
        return not children

    def _unwrap(l: _LeafWrapper):
      assert isinstance(l, _LeafWrapper)
      return l.wrapped_leaf

    # This logic is based on equinox.tree_at, but kept separate to avoid
    # depending on equinox.
    def _process_one(node, multiple=multiple):