def apply(
      self,
      fn: Callable[..., Any],
      *,
      with_keypath: bool = False,
      keep_selected: bool = False,
  ) -> Any:
    """Replaces each selected node with the result of applying this function.

    Args:
      fn: Function to apply to each selected node. This function should take a
        PyTree (if with_keypath=False) or a KeyPath and a PyTree (if
        with_keypath=True) and return a replacement PyTree.
      with_keypath: Whether to pass the keypath as the first argument to the
        callable.
      keep_selected: Whether to keep the nodes selected. If True, returns the
        modified selection; if False, rebuilds the tree after replacing.

    Returns:
      Either a modified `Selection` (if keep_selected=True) or a rebuilt version
      of the original tree, each with the replacements applied.
    """
    if with_keypath:
      new_values = collections.OrderedDict(
          [(k, fn(k, v)) for k, v in self.selected_by_path.items()]
      )
    else:
      new_values = collections.OrderedDict(
          [(k, fn(v)) for k, v in self.selected_by_path.items()]
      )
    new_selection = Selection(
        selected_by_path=new_values, remainder=self.remainder
    )
    if keep_selected:
      return new_selection
    else:
      return new_selection.deselect()
Replaces each selected node with the result of applying this function.

    Args:
      fn: Function to apply to each selected node. This function should take a
        PyTree (if with_keypath=False) or a KeyPath and a PyTree (if
        with_keypath=True) and return a replacement PyTree.
      with_keypath: Whether to pass the keypath as the first argument to the
        callable.
      keep_selected: Whether to keep the nodes selected. If True, returns the
        modified selection; if False, rebuilds the tree after replacing.

    Returns:
      Either a modified `Selection` (if keep_selected=True) or a rebuilt version
      of the original tree, each with the replacements applied.
    
penzai/core/selectors.py