def inline_groups(
    tree: Any,
    parent_filter: Callable[[Any], bool],
    child_filter: Callable[[Any], bool],
) -> Any:
  """Inlines sequential nodes into their parents if possible.

  This function finds nodes that match ``child_filter`` within nodes that match
  ``parent_filter``, and splices the sublayers of the child node into the
  parent, removing the child. This can be used to flatten a nested structure of
  `Sequential` or `NamedGroup` objects into a new structer with a smaller depth.

  The logic applies recursively: if a node matches both the parent and the
  child filter, it may be inlined into its parent after its sublayers are
  inlined into it.

  For the common case where you wish to inline "anonymous" groups (instances
  of type `Sequential` but not a more specific subclass of `Sequential`), you
  can use the convenience wrapper `inline_anonymous_sequentials`.

  Args:
    tree: The tree to process.
    parent_filter: A function that returns True on the nodes that we want to
      inline sublayers into.
    child_filter: A function that returns True on the nodes that we want to
      remove and replace with the inlined sequence of its sublayers.

  Returns:
    A copy of ``tree`` that inlines nodes that match ``child_filter`` into their
    parents whenever their parents match ``parent_filter``, as long as they
    are subclasses of `Sequential`, `CheckedSequential`, or `NamedGroup`.
  """

  def _step(subtree):
    # Process children.
    with_processed_children = (
        selectors.select(subtree).at_children().apply(_step)
    )
    # Check for inlining opportunities at this level.
    if is_sequential_or_named(with_processed_children) and parent_filter(
        with_processed_children
    ):
      new_sublayers = []
      for sublayer in with_processed_children.sublayers:
        if is_sequential_or_named(sublayer) and child_filter(sublayer):
          # Inline the child's children into the parent's children.
          new_sublayers.extend(sublayer.sublayers)
        else:
          new_sublayers.append(sublayer)

      # Substitute the new flattened children.
      return dataclasses.replace(
          with_processed_children, sublayers=new_sublayers
      )
    else:
      return with_processed_children

  return _step(tree)
Inlines sequential nodes into their parents if possible.

  This function finds nodes that match ``child_filter`` within nodes that match
  ``parent_filter``, and splices the sublayers of the child node into the
  parent, removing the child. This can be used to flatten a nested structure of
  `Sequential` or `NamedGroup` objects into a new structer with a smaller depth.

  The logic applies recursively: if a node matches both the parent and the
  child filter, it may be inlined into its parent after its sublayers are
  inlined into it.

  For the common case where you wish to inline "anonymous" groups (instances
  of type `Sequential` but not a more specific subclass of `Sequential`), you
  can use the convenience wrapper `inline_anonymous_sequentials`.

  Args:
    tree: The tree to process.
    parent_filter: A function that returns True on the nodes that we want to
      inline sublayers into.
    child_filter: A function that returns True on the nodes that we want to
      remove and replace with the inlined sequence of its sublayers.

  Returns:
    A copy of ``tree`` that inlines nodes that match ``child_filter`` into their
    parents whenever their parents match ``parent_filter``, as long as they
    are subclasses of `Sequential`, `CheckedSequential`, or `NamedGroup`.
  
penzai/nn/grouping.py