inline_groups

penzai.nn.grouping.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.

Parameters:

• 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.