def freeze_local_states(
    handled: WithFunctionalLocalState, states: dict[str, Any]
) -> Any:
  """Embeds the given states into a handled model, and removes the handler.

  This function converts a functionalized stateful model into a single effectful
  one, making it easier to inspect, patch, or serialize. It is roughly the
  inverse of `handle_local_states`, and `handle_local_states` can be used
  to re-install the handler.

  Args:
    handled: A WithFunctionalLocalState wrapping a model, as constructed using
      `handle_local_states`.
    states: A state dictionary, either from `handle_local_states` or from
      calling the handled model to get updated states.

  Returns:
    A copy of the original model, without the `WithFunctionalLocalState`
    wrapper, and with each of the states embedded as a `FrozenLocalStateRequest`
    inside the tree.
  """
  already_made_request_for_state = set()

  def _hole_for_ref(ref: HandledLocalStateRef):
    if ref.was_explicitly_named:
      name = ref.name
      if name in already_made_request_for_state:
        return SharedLocalStateRequest(name=name, category=ref.category)
      else:
        already_made_request_for_state.add(name)
    else:
      name = None
    return FrozenLocalStateRequest(
        state=states[ref.name], name=name, category=ref.category
    )

  return (
      selectors.select(handled.body)
      .at_instances_of(effect_base.HandledEffectRef)
      .where(lambda ref: ref.handler_id == handled.handler_id)
      .apply(_hole_for_ref)
  )
Embeds the given states into a handled model, and removes the handler.

  This function converts a functionalized stateful model into a single effectful
  one, making it easier to inspect, patch, or serialize. It is roughly the
  inverse of `handle_local_states`, and `handle_local_states` can be used
  to re-install the handler.

  Args:
    handled: A WithFunctionalLocalState wrapping a model, as constructed using
      `handle_local_states`.
    states: A state dictionary, either from `handle_local_states` or from
      calling the handled model to get updated states.

  Returns:
    A copy of the original model, without the `WithFunctionalLocalState`
    wrapper, and with each of the states embedded as a `FrozenLocalStateRequest`
    inside the tree.
  
penzai/deprecated/v1/data_effects/local_state.py