IsolatedSubmodel

class penzai.toolshed.isolate_submodel.IsolatedSubmodel

Bases: Struct

An isolated part of a submodel, with its saved inputs and outputs.

When initially populated by isolate_submodel, any states used by the submodel will be functionalized and treated as part of the saved input and output, and any shared parameters, side inputs, or random streams will be captured and stored in submodel as well.

Variables:

• submodel (pz.Layer) – An individual layer from the larger model.

• saved_input (Any) – Input that was passed to the submodel when we isolated it.

• saved_output (Any) – Saved output that the submodel should produce when called with saved_input.

Methods

__init__(submodel, saved_input, saved_output)

Attributes

submodel
saved_input
saved_output

Inherited Methods

(expand to view inherited methods)

attributes_dict()	Constructs a dictionary with all of the fields in the class.
from_attributes(**field_values)	Directly instantiates a struct given all of its fields.
key_for_field(field_name)	Generates a JAX PyTree key for a given field name.
select()	Wraps this struct in a selection, enabling functional-style mutations.
tree_flatten()	Flattens this tree node.
tree_flatten_with_keys()	Flattens this tree node with keys.
tree_unflatten(aux_data, children)	Unflattens this tree node.
treescope_color()	Computes a CSS color to display for this object in treescope.