call_and_extract_submodel

penzai.toolshed.isolate_submodel.call_and_extract_submodel(submodel_selection: pz.Selection[pz.Layer], argument: Any) → 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 large 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.Layer within a larger model (also a pz.Layer). Must contain exactly one selected subtree, and the selected subtree must be a layer.

• argument – Argument to call the full model with.

Returns:

An isolated view of the selected submodel, with the inputs, outputs, states and any shared parameters 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.