call_and_extract_submodel

penzai.toolshed.isolate_submodel.call_and_extract_submodel(submodel_selection: pz.Selection[pz.nn.Layer], argument: Any, /, **side_inputs) → IsolatedSubmodel | list[IsolatedSubmodel]

Calls a model with an argument, and captures the selected submodel.

This function is designed to enable quickly exploring the behavior of small parts of a larger model, and allows reproducing the in-context behavior of the subpart without actually having to run the full model.

Parameters:

• submodel_selection – A selection of a single pz.nn.Layer within a larger model (also a pz.nn.Layer). Must contain exactly one selected subtree, and the selected subtree must be a layer.

• argument – Argument to call the full model with.

• **side_inputs – Side inputs for the model.

Returns:

An isolated view of the selected submodel, with the inputs, outputs, and states captured so that the submodel can be analyzed without invoking the larger model. If the submodel was called multiple times, each call will be captured separately and returned as a list. Note that all parameters will also be frozen in the result.