def apply_and_inline(
      self,
      fn: Callable[..., Iterable[Any]],
      *,
      with_keypath: bool = False,
  ) -> Any:
    """Replaces selected list/tuple items with a sequence of new items.

    This function should only be called when the selected elements are all
    children of lists or tuples. It removes the original selected element from
    the sequence, then inserts the results of calling ``fn`` at that location.
    This is similar to the "flatmap" operation in some other languages.

    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``). It should return a replacement iterable of
        PyTrees, each of which will be inserted into the parent container in
        place of the selected node.
      with_keypath: Whether to pass the keypath as the first argument to the
        callable.

    Returns:
      A replacement of the original tree with the selected elements replaced
      with the outputs of ``fn``.

    Raises:
      ValueError: If any selected node is not the child of a list or tuple.
    """
    error_extra_if_callable = (
        "\n\nIf you are trying to replace a layer (or other callable) with a"
        " sequence of callables, consider"
        " replacing it with a Sequential layer instead, using something like"
        " `.apply(lambda x: Sequential(...))` instead of"
        " `.apply_and_inline(lambda x: ...)`."
    )

    def process_subtree(subtree, keypath, parent):
      if isinstance(subtree, SelectionHole):
        # Holes inside of lists or tuples will be processed by their parent.
        # If we get here, that's an error.
        if callable(subtree):
          # Give some more help if the subtree is callable, since that might
          # mean they are trying to insert logic before/after this node.
          error_extra = error_extra_if_callable
        else:
          error_extra = ""
        if parent is None:
          raise ValueError(
              "apply_and_inline requires all selected nodes to be children of"
              " lists or tuples, but the root node was selected."
              + error_extra
          )
        else:
          raise ValueError(
              "apply_and_inline requires all selected nodes to be children of"
              " lists or tuples, but the node at keypath"
              f" {jax.tree_util.keystr(keypath)} was contained in a node of"
              f" type {type(parent).__name__}"
              + error_extra
          )
      elif isinstance(subtree, SelectionQuote):
        # We are resolving the selection, so unquote here.
        return subtree.quoted
      elif isinstance(subtree, (list, tuple)):
        # Process and inline any selected children.
        new_children = []
        for i, child in enumerate(subtree):
          child_keypath = keypath + (jax.tree_util.SequenceKey(i),)
          if isinstance(child, SelectionHole):
            child_value = self.selected_by_path[child.path]
            if with_keypath:
              replacements = fn(child_keypath, child_value)
            else:
              replacements = fn(child_value)
            new_children.extend(replacements)
          else:
            # Not a selected node; process recursively.
            new_children.append(
                process_subtree(child, keypath=child_keypath, parent=subtree)
            )
        return type(subtree)(new_children)
      else:
        # Recursively process children.
        maybe_children = penzai_tree_util.tree_flatten_exactly_one_level(
            subtree
        )
        if maybe_children is None:
          return subtree
        else:
          children_by_key, treedef = maybe_children
          return treedef.unflatten([
              process_subtree(child, keypath + (key,), subtree)
              for key, child in children_by_key
          ])

    return process_subtree(self.remainder, (), None)
Replaces selected list/tuple items with a sequence of new items.

    This function should only be called when the selected elements are all
    children of lists or tuples. It removes the original selected element from
    the sequence, then inserts the results of calling ``fn`` at that location.
    This is similar to the "flatmap" operation in some other languages.

    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``). It should return a replacement iterable of
        PyTrees, each of which will be inserted into the parent container in
        place of the selected node.
      with_keypath: Whether to pass the keypath as the first argument to the
        callable.

    Returns:
      A replacement of the original tree with the selected elements replaced
      with the outputs of ``fn``.

    Raises:
      ValueError: If any selected node is not the child of a list or tuple.
    
penzai/core/selectors.py