def saving_all_intermediates(
    model: pz.nn.Layer,
    intermediates_group: str | None = "intermediates",
    shapes_only: bool = False,
) -> pz.nn.Layer:
  """Returns a copy of `model` with `SaveIntermediate` layers inserted.

  Args:
    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.
  """
  if shapes_only:
    inserted_cls = SaveIntermediateShape
  else:
    inserted_cls = SaveIntermediate

  # Recursively traverse the model to identify all of the places where
  # we might need to add shape annotations, and add a inserted_cls
  # layer there. We try to only save each intermediate value once.
  def add_intercept_values_inside(cur_layer: pz.nn.Layer) -> pz.nn.Layer:
    """Inserts inserted_cls layers around sublayers of a layer."""
    if isinstance(
        cur_layer, pz.nn.Sequential | pz.nn.NamedGroup | pz.nn.CheckedSequential
    ):
      sublayers = cur_layer.sublayers
      new_sublayers = []
      for i, sublayer in enumerate(sublayers):
        new_sublayers.append(add_intercept_values_inside(sublayer))
        if not isinstance(sublayer, pz.nn.Identity) and i != len(sublayers) - 1:
          new_sublayers.append(inserted_cls())
      return dataclasses.replace(cur_layer, sublayers=new_sublayers)
    else:
      return (
          pz.select(cur_layer)
          .at_children()
          .at_instances_of(pz.nn.Layer)
          .apply(add_intercept_values_around)
      )

  def add_intercept_values_around(cur_layer: pz.nn.Layer) -> pz.nn.Layer:
    """Inserts inserted_cls layers around a layer and inside sublayers."""
    if isinstance(
        cur_layer, pz.nn.Sequential | pz.nn.NamedGroup | pz.nn.CheckedSequential
    ):
      sublayers = cur_layer.sublayers
      new_sublayers = [inserted_cls()]
      for sublayer in sublayers:
        new_sublayers.append(add_intercept_values_inside(sublayer))
        if not isinstance(sublayer, pz.nn.Identity):
          new_sublayers.append(inserted_cls())
      return dataclasses.replace(cur_layer, sublayers=new_sublayers)
    elif isinstance(cur_layer, pz.nn.Identity):
      # Add intermediates before the identity but not after (since it won't
      # change them).
      return pz.nn.Sequential([
          inserted_cls(),
          add_intercept_values_inside(cur_layer),
      ])
    else:
      # Make a new sequential so we can add intermediates around it.
      return pz.nn.Sequential([
          inserted_cls(),
          add_intercept_values_inside(cur_layer),
          inserted_cls(),
      ])

  if intermediates_group is None:
    return add_intercept_values_around(model)
  else:
    with pz.scoped_auto_state_var_labels(intermediates_group):
      return add_intercept_values_around(model)
Returns a copy of `model` with `SaveIntermediate` layers inserted.

  Args:
    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.
  
penzai/toolshed/save_intermediates.py