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)
      )
Inserts inserted_cls layers around sublayers of a layer.
penzai/toolshed/save_intermediates.py