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
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)
          
penzai/core/selectors.py