hoist_constant_side_inputs

penzai.data_effects.side_input.hoist_constant_side_inputs(tree: Any, hoist_predicate: Callable[[Tag, Any], bool] = <function <lambda>>, scoped: bool = True) → tuple[Any, dict[Tag, Any]]

Extracts all constant side inputs from a tree so they can be re-handled.

This function finds all WithConstantSideInputs handlers in a tree, and extracts the side inputs from those handlers, and replaces the bound side input references with unhandled requests. You can then manipulate the tree and finally re-handle the side inputs using WithConstantSideInputs.handling.

This function is useful when a tree contains constant side inputs inside a subtree, and you want to pull out an even smaller subtree that uses those side inputs, or move around subtrees that use side inputs. A particular instance of this is shared parameters, which can be represented in Penzai as side inputs with a constant value provided by a WithConstantSideInputs handler. If you want to extract a layer that uses shared parameters, you can first hoist out the constant values of all shared parameters, extract the submodel you want, and then re-bind the shared parameters of that subtree alone.

Calling hoist_constant_side_inputs on a layer and then immediately calling WithConstantSideInputs.handling on the results will usually produce a layer with the same observable behavior, although the tree structure may differ. However, if scoped is False, it is possible that this function will raise an error due to tag conflicts.

Parameters:

• tree – A tree to hoist constant side inputs from. Usually will contain WithConstantSideInputs handlers somewhere inside it.

• hoist_predicate – Optional predicate to determine which side inputs to try to hoist, given the tag and value. Defaults to hoisting all side inputs.

• scoped – Whether to modify the tags for the hoisted side inputs so that they remain unique even if different side inputs use different tags.

Returns:

Tuple (effectful_tree, hoisted_values) where effectful_tree is a copy of tree with all hoisted side inputs replaced by SideInputRequest instances and all WithConstantSideInputs removed or modified to not include the hoisted side inputs, and hoisted_values are the values that were bound by those constant side inputs, as needed by WithConstantSideInputs.handling.