def at_subtrees_where(
      self,
      filter_fn: Callable[..., bool],
      *,
      with_keypath: bool = False,
      absolute_keypath: bool = False,
      innermost: bool = False,
  ) -> "Selection":
    """Selects subtrees of selected nodes where a function evaluates to True.

    Note that a selection cannot contain a node that is the descendant of
    another selected node. If ``innermost=False``, we return the outermost
    node, whereas if ``innermost=True`` we return the innermost.

    If you want to apply a modification to *all* matches of a function, even
    if they are nested, you can use a pattern like ::

      selection = select(value)
      while not selection.empty():
        selection = selection.at_subtrees_where(foo).apply(
            bar, keep_selected=True)
      new_value = selection.deselect()

    More complex modifications can also be made using a manual traversal, e.g.::

      def traverse(subtree):
        # ... process the subtree before recursive call ...
        subtree = select(subtree).at_children().apply(traverse)
        # ... process the subtree after the recursive call ...
        return subtree

      new_value = traverse(value)

    Args:
      filter_fn: A function determining which subtrees to select. Should be
        deterministic, and may be called more than once. This function should
        take a PyTree (if ``with_keypath=False``) or a KeyPath and a PyTree (if
        ``with_keypath=True``).
      with_keypath: Whether to pass a keypath as the first argument to the
        callable.
      absolute_keypath: Whether to pass the keypath relative to the root of the
        original tree (if True) or the keypath relative to the currently
        selected node (if False). Ignored if ``with_keypath`` is False.
      innermost: Whether to select the innermost subtree(s) for which the filter
        function is true, instead of the first subtrees encountered.

    Returns:
      A new selection that selects the desired subtrees.
    """

    def safe_filter_fn(*args):
      result = filter_fn(*args)
      if not isinstance(result, bool):
        raise TypeError(
            "Filter function for at_subtrees_where must return a bool, not"
            f" {type(result)}"
        )
      return result

    with _wrap_selection_errors(self):
      if with_keypath or innermost:
        if with_keypath:
          wrapped_filter_fn = safe_filter_fn
        else:
          wrapped_filter_fn = lambda _, s: safe_filter_fn(s)

        def process_subtree(keypath, leaf_or_subtree) -> tuple[bool, Any]:
          """Recursively walks subtrees one level at a time.

          Args:
            keypath: Keypath to this subtree
            leaf_or_subtree: The subtree (or a leaf)

          Returns:
            (found_any, processed_pytree)
          """
          if not innermost:
            found_here = wrapped_filter_fn(keypath, leaf_or_subtree)
            if found_here:
              # Select it.
              return True, _InProgressSelectionBoundary(leaf_or_subtree)

          maybe_children = penzai_tree_util.tree_flatten_exactly_one_level(
              leaf_or_subtree
          )
          if maybe_children:
            # Recursive step.
            keyed_children, treedef = maybe_children
            new_children = []
            any_descendant_selected = False
            for key, child in keyed_children:
              found_in_child, new_child = process_subtree(
                  keypath + (key,), child
              )
              new_children.append(new_child)
              any_descendant_selected = (
                  any_descendant_selected or found_in_child
              )

            if not any_descendant_selected and wrapped_filter_fn(
                keypath, leaf_or_subtree
            ):
              # No descendant was selected, so select this tree.
              return True, _InProgressSelectionBoundary(leaf_or_subtree)
            elif any_descendant_selected:
              # A descendant was selected, use the descendant.
              return True, treedef.unflatten(new_children)
            else:
              # Didn't find anything.
              return False, leaf_or_subtree
          else:
            # This is a leaf.
            return False, leaf_or_subtree

        def process_selected(keypath, selected_subtree):
          """Processes one of the selected subtrees."""
          if absolute_keypath:
            # Start with this selection's keypath
            _, result = process_subtree(keypath, selected_subtree)
            return result
          else:
            # Start with the empty keypath (relative to this selection)
            _, result = process_subtree((), selected_subtree)
            return result

        return _build_selection_from_boundary(
            self.apply(process_selected, with_keypath=True)
        )

      else:

        def process_leaf_or_filtered(leaf_or_subtree):
          """Processes leaves or values where filter_fn is True."""
          if safe_filter_fn(leaf_or_subtree):
            return _InProgressSelectionBoundary(leaf_or_subtree)
          else:
            return leaf_or_subtree

        def process_selected(selected_subtree):
          """Processes one of the selected subtrees."""
          return jax.tree_util.tree_map(
              process_leaf_or_filtered, selected_subtree, is_leaf=safe_filter_fn
          )

        return _build_selection_from_boundary(self.apply(process_selected))
Selects subtrees of selected nodes where a function evaluates to True.

    Note that a selection cannot contain a node that is the descendant of
    another selected node. If ``innermost=False``, we return the outermost
    node, whereas if ``innermost=True`` we return the innermost.

    If you want to apply a modification to *all* matches of a function, even
    if they are nested, you can use a pattern like ::

      selection = select(value)
      while not selection.empty():
        selection = selection.at_subtrees_where(foo).apply(
            bar, keep_selected=True)
      new_value = selection.deselect()

    More complex modifications can also be made using a manual traversal, e.g.::

      def traverse(subtree):
        # ... process the subtree before recursive call ...
        subtree = select(subtree).at_children().apply(traverse)
        # ... process the subtree after the recursive call ...
        return subtree

      new_value = traverse(value)

    Args:
      filter_fn: A function determining which subtrees to select. Should be
        deterministic, and may be called more than once. This function should
        take a PyTree (if ``with_keypath=False``) or a KeyPath and a PyTree (if
        ``with_keypath=True``).
      with_keypath: Whether to pass a keypath as the first argument to the
        callable.
      absolute_keypath: Whether to pass the keypath relative to the root of the
        original tree (if True) or the keypath relative to the currently
        selected node (if False). Ignored if ``with_keypath`` is False.
      innermost: Whether to select the innermost subtree(s) for which the filter
        function is true, instead of the first subtrees encountered.

    Returns:
      A new selection that selects the desired subtrees.
    
penzai/core/selectors.py