def hoist_shared_state_requests(
    tree: Any,
    unsafe: bool = False,
) -> tuple[Any, dict[str, InitialLocalStateRequest | FrozenLocalStateRequest]]:
  r"""Hoists out the value of shared states in a pytree.

  This function is a helper for manipulating Penzai models that contain shared
  states. Ordinarily, shared states in a Penzai model are represented as some
  combination of:

  * exactly one of:

    * a single `FrozenLocalStateRequest` with a value

    * multiple `InitialLocalStateRequest` nodes with identical initializers

  * followed by one or more `SharedLocalStateRequest` nodes

  where all such requests have the same explicit name. This is convenient for
  manipulating the model as a whole, but can make it somewhat annoying to
  extract a small part of a model that uses a shared state defined elsewhere.

  This function takes a tree of this form and returns a new tree that only
  contains a `SharedLocalStateRequest` whenever there is a state that is used in
  multiple places, along with a dictionary mapping each state name to the
  concrete definition it uses. The new tree can be freely manipulated, and then
  a single copy of the state can be re-embedded using `embed_shared_states`.

  Args:
    tree: A tree where each shared variable appears either once as a
      `FrozenLocalStateRequest` or multiple times as identical
      `InitialLocalStateRequest` nodes, followed by some number of
      `SharedLocalStateRequest` nodes.
    unsafe: If True, `tree` can have multiple `FrozenLocalStateRequest` or
      `InitialLocalStateRequest` nodes with different initializers, and one will
      be picked arbitrarily.

  Returns:
    A tuple of ``(new_tree, state_defs)``, where ``new_tree`` is a copy of
    ``tree`` where all shared states have been replaced by a
    `SharedLocalStateRequest`, and ``state_defs`` is a dictionary mapping each
    shared state name to the corresponding `InitialLocalStateRequest` or
    `FrozenLocalStateRequest`. These can be passed to
    `embed_shared_state_requests` to rebuild the original tree.
  """
  state_counter = collections.Counter(
      req.name
      for req in (
          selectors.select(tree)
          .at_instances_of(
              InitialLocalStateRequest
              | FrozenLocalStateRequest
              | SharedLocalStateRequest
          )
          .get_sequence()
      )
  )

  hoisted_vals = {}

  def _hoist_it(request: InitialLocalStateRequest | FrozenLocalStateRequest):
    if request.name not in hoisted_vals:
      hoisted_vals[request.name] = request
    elif (
        isinstance(request, InitialLocalStateRequest)
        and isinstance(hoisted_vals[request.name], InitialLocalStateRequest)
        and request.state_initializer
        is hoisted_vals[request.name].state_initializer
    ):
      pass
    elif not unsafe:
      raise ValueError(
          "Detected two local states with the same explicit name"
          f" {repr(request.name)} but different initializers or values! This is"
          " only allowed `unsafe` is set to `True`."
      )
    return SharedLocalStateRequest(name=request.name, category=request.category)

  hoisted_tree = (
      selectors.select(tree)
      .at_instances_of(InitialLocalStateRequest | FrozenLocalStateRequest)
      .where(
          lambda request: (
              request.name is not None and state_counter[request.name] > 1
          )
      )
      .apply(_hoist_it)
  )
  return hoisted_tree, hoisted_vals
Hoists out the value of shared states in a pytree.

  This function is a helper for manipulating Penzai models that contain shared
  states. Ordinarily, shared states in a Penzai model are represented as some
  combination of:

  * exactly one of:

    * a single `FrozenLocalStateRequest` with a value

    * multiple `InitialLocalStateRequest` nodes with identical initializers

  * followed by one or more `SharedLocalStateRequest` nodes

  where all such requests have the same explicit name. This is convenient for
  manipulating the model as a whole, but can make it somewhat annoying to
  extract a small part of a model that uses a shared state defined elsewhere.

  This function takes a tree of this form and returns a new tree that only
  contains a `SharedLocalStateRequest` whenever there is a state that is used in
  multiple places, along with a dictionary mapping each state name to the
  concrete definition it uses. The new tree can be freely manipulated, and then
  a single copy of the state can be re-embedded using `embed_shared_states`.

  Args:
    tree: A tree where each shared variable appears either once as a
      `FrozenLocalStateRequest` or multiple times as identical
      `InitialLocalStateRequest` nodes, followed by some number of
      `SharedLocalStateRequest` nodes.
    unsafe: If True, `tree` can have multiple `FrozenLocalStateRequest` or
      `InitialLocalStateRequest` nodes with different initializers, and one will
      be picked arbitrarily.

  Returns:
    A tuple of ``(new_tree, state_defs)``, where ``new_tree`` is a copy of
    ``tree`` where all shared states have been replaced by a
    `SharedLocalStateRequest`, and ``state_defs`` is a dictionary mapping each
    shared state name to the corresponding `InitialLocalStateRequest` or
    `FrozenLocalStateRequest`. These can be passed to
    `embed_shared_state_requests` to rebuild the original tree.
  
penzai/deprecated/v1/data_effects/local_state.py