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(),
      ])
Inserts inserted_cls layers around a layer and inside sublayers.
penzai/toolshed/save_intermediates.py