def add_intercept_values_inside(cur_layer: pz.Layer) -> pz.Layer:
    """Inserts _InterceptValues 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 sublayer in sublayers:
        new_sublayers.append(add_intercept_values_inside(sublayer))
        if not isinstance(sublayer, pz.nn.Identity):
          new_sublayers.append(_InterceptValues.build())
      if new_sublayers and isinstance(new_sublayers[-1], _InterceptValues):
        new_sublayers.pop()
      return dataclasses.replace(cur_layer, sublayers=new_sublayers)
    else:
      return (
          pz.select(cur_layer)
          .at_children()
          .at_instances_of(pz.Layer)
          .apply(add_intercept_values_around)
      )
Inserts _InterceptValues layers around sublayers of a layer.
penzai/deprecated/v1/toolshed/interleave_intermediates.py