def call_and_extract_submodel(
    submodel_selection: pz.Selection[pz.nn.Layer],
    argument: Any,
    /,
    **side_inputs,
) -> IsolatedSubmodel | list[IsolatedSubmodel]:
  """Calls a model with an argument, and captures the selected submodel.

  This function is designed to enable quickly exploring the behavior of small
  parts of a larger model, and allows reproducing the in-context behavior of the
  subpart without actually having to run the full model.

  Args:
    submodel_selection: A selection of a single `pz.nn.Layer` within a larger
      model (also a `pz.nn.Layer`). Must contain exactly one selected subtree,
      and the selected subtree must be a layer.
    argument: Argument to call the full model with.
    **side_inputs: Side inputs for the model.

  Returns:
    An isolated view of the selected submodel, with the inputs, outputs, and
    states captured so that the submodel can be analyzed without invoking the
    larger model. If the submodel was called multiple times, each call will be
    captured separately and returned as a list. Note that all parameters will
    also be frozen in the result.
  """
  # Check the selection.
  if submodel_selection.count() != 1:
    raise ValueError("submodel_selection must contain exactly one sublayer.")
  selected_layer = submodel_selection.get()
  if not isinstance(selected_layer, pz.nn.Layer):
    raise ValueError("submodel_selection must select a layer.")
  if not isinstance(submodel_selection.deselect(), pz.nn.Layer):
    raise ValueError(
        "The tree that submodule_selection was built from must be a layer."
    )

  capturer = IsolationCapturer(
      wrapped=submodel_selection.get(), saved_calls=pz.StateVariable([])
  )
  model_with_capturer = submodel_selection.set(capturer)

  # Call it to extract the necessary parts.
  _ = model_with_capturer(argument, **side_inputs)

  # Return the result.
  if len(capturer.saved_calls.value) == 1:
    return capturer.saved_calls.value[0]
  else:
    return capturer.saved_calls.value
Calls a model with an argument, and captures the selected submodel.

  This function is designed to enable quickly exploring the behavior of small
  parts of a larger model, and allows reproducing the in-context behavior of the
  subpart without actually having to run the full model.

  Args:
    submodel_selection: A selection of a single `pz.nn.Layer` within a larger
      model (also a `pz.nn.Layer`). Must contain exactly one selected subtree,
      and the selected subtree must be a layer.
    argument: Argument to call the full model with.
    **side_inputs: Side inputs for the model.

  Returns:
    An isolated view of the selected submodel, with the inputs, outputs, and
    states captured so that the submodel can be analyzed without invoking the
    larger model. If the submodel was called multiple times, each call will be
    captured separately and returned as a list. Note that all parameters will
    also be frozen in the result.
  
penzai/toolshed/isolate_submodel.py