def set_by_path(
      self, replacer: Mapping[KeyPath, Any] | Callable[[KeyPath], Any]
  ) -> Any:
    """Replaces the selected subtree(s) based on their path(s).

    If you need both the value and the key, see
    ``.apply(fn, with_keypath=True)``.

    Args:
      replacer: A mapping from key paths to replacements, or a function that
        builds such a mapping. Passing ``self.selection_by_path`` will return
        the original tree unchanged.

    Returns:
      A modified version of the original tree, with replacements taken from the
      replacer.
    """
    if callable(replacer):
      replacer = {k: replacer(k) for k in self.selected_by_path.keys()}
    return self.apply(lambda k, v: replacer[k], with_keypath=True)
Replaces the selected subtree(s) based on their path(s).

    If you need both the value and the key, see
    ``.apply(fn, with_keypath=True)``.

    Args:
      replacer: A mapping from key paths to replacements, or a function that
        builds such a mapping. Passing ``self.selection_by_path`` will return
        the original tree unchanged.

    Returns:
      A modified version of the original tree, with replacements taken from the
      replacer.
    
penzai/core/selectors.py