saving_all_intermediates

saving_all_intermediates#

penzai.toolshed.save_intermediates.saving_all_intermediates(model: pz.nn.Layer, intermediates_group: str | None = 'intermediates', shapes_only: bool = False) → pz.nn.Layer[source]#

Returns a copy of model with SaveIntermediate layers inserted.

Parameters:

  • model – The model to copy.

  • intermediates_group – The name of the group to use for the intermediate value Variable labels. If None, the intermediates will be given unique labels.

  • shapes_only – If True, only store the shapes of the intermediate values, not the values themselves.

Returns:

A copy of the layer or model with empty SaveIntermediate layers added before and after every sublayer in the model.

Contents