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
          )
Processes one of the selected subtrees.
penzai/core/selectors.py