def add_intercept_shapes_inside(cur_layer: pz.Layer) -> pz.Layer:
    """Inserts `_InterceptShapes` layers around all sublayers of a layer."""
    if isinstance(
        cur_layer, pz.nn.Sequential | pz.nn.NamedGroup | pz.nn.CheckedSequential
    ):
      sublayers = cur_layer.sublayers
      new_sublayers = [_InterceptShapes.build()]
      for sublayer in sublayers:
        new_sublayers.append(add_intercept_shapes_inside(sublayer))
        new_sublayers.append(_InterceptShapes.build())
      return dataclasses.replace(cur_layer, sublayers=new_sublayers)
    else:
      return (
          pz.select(cur_layer)
          .at_children()
          .at_instances_of(pz.Layer)
          .apply(add_intercept_shapes_around)
      )
Inserts `_InterceptShapes` layers around all sublayers of a layer.
penzai/deprecated/v1/toolshed/annotate_shapes.py