def where(
      self: "Selection[SelectedSubtree]",
      filter_fn: Callable[..., bool],
      *,
      with_keypath: bool = False,
  ) -> "Selection[SelectedSubtree]":
    """Filters to only a subset of selected nodes based on a condition.

    Args:
      filter_fn: Function to determine whether to keep a node in the selection.
        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 the keypath as the first argument to the
        callable.

    Returns:
      A new selection that includes only the selected parts where ``filter_fn``
      evaluates to true.
    """
    with _wrap_selection_errors(self):
      keep = _InProgressSelectionBoundary
      if with_keypath:
        new_with_boundary = self.apply(
            lambda k, v: keep(v) if filter_fn(k, v) else v, with_keypath=True
        )
      else:
        new_with_boundary = self.apply(lambda v: keep(v) if filter_fn(v) else v)

      return _build_selection_from_boundary(new_with_boundary)
Filters to only a subset of selected nodes based on a condition.

    Args:
      filter_fn: Function to determine whether to keep a node in the selection.
        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 the keypath as the first argument to the
        callable.

    Returns:
      A new selection that includes only the selected parts where ``filter_fn``
      evaluates to true.
    
penzai/core/selectors.py