unwrap (pmrf.unwrap)

pmrf.unwrap(tree: AbstractUnwrappable[T] | T, only_if: Callable[[Any], bool] = None, cascade: bool = True) → T

Recursively resolves AbstractUnwrappable nodes within a PyTree.

By default, unwrapping is performed inside-out (bottom-up) across the entire tree. Every AbstractUnwrappable node is replaced by the result of its unwrap() method.

If the only_if predicate is provided, unwrapping is conditionally gated. The cascade parameter controls how descendants of matching nodes are handled.

Behavior with only_if:

  1. If cascade=True (default): The tree is searched top-down. Once a node satisfies only_if, that node and ALL of its AbstractUnwrappable descendants are fully resolved.

  2. If cascade=False: The tree is traversed bottom-up. Unwrapping ONLY triggers for specific nodes that satisfy the condition. Unmatching descendants are left wrapped.

Args:

tree: The PyTree to unwrap. only_if: An optional predicate function Callable[[Any], bool]. cascade: If True, unwrapping cascades to all descendants of a matched node.

If False, only nodes strictly evaluating to True are unwrapped.

Returns:

A new PyTree with the appropriate AbstractUnwrappable nodes resolved.