def embed_shared_state_requests(
    tree: Any,
    state_requests: dict[
        str, InitialLocalStateRequest | FrozenLocalStateRequest
    ],
) -> Any:
  """Embeds shared state requests into a tree.

  This function is the inverse of hoist_shared_state_requests, and can be used
  to re-embed the initial value of explicitly-named shared states into a tree
  so that they can be handled using an ordinary handler.

  The intended use of `embed_shared_state_requests` is when you want to extract
  parts of a model with shared states, or embed a model with shared states into
  a larger model. This function ensures that there is only one initial value
  for each named state, regardless of how many times it is used.

  Args:
    tree: A tree where each shared variable appears as a
      `SharedLocalStateRequest`.
    state_requests: A dictionary mapping state names to the corresponding
      `InitialLocalStateRequest` or `FrozenLocalStateRequest`.

  Returns:
    A tree where the first appearance of each shared state is replaced by the
    corresponding `InitialLocalStateRequest` or `FrozenLocalStateRequest`, and
    all other appearances are kept as `SharedLocalStateRequest`.
  """
  pending = dict(state_requests)

  def _maybe_embed_it(request: SharedLocalStateRequest):
    if request.name in pending:
      replacement = pending.pop(request.name)
      if request.category != replacement.category:
        raise ValueError(
            "Tried to substitute values with mismatched categories! Got"
            f" {replacement.category}, expected {request.category}"
        )

      return replacement
    else:
      return request

  return (
      selectors.select(tree)
      .at_instances_of(SharedLocalStateRequest)
      .apply(_maybe_embed_it)
  )
Embeds shared state requests into a tree.

  This function is the inverse of hoist_shared_state_requests, and can be used
  to re-embed the initial value of explicitly-named shared states into a tree
  so that they can be handled using an ordinary handler.

  The intended use of `embed_shared_state_requests` is when you want to extract
  parts of a model with shared states, or embed a model with shared states into
  a larger model. This function ensures that there is only one initial value
  for each named state, regardless of how many times it is used.

  Args:
    tree: A tree where each shared variable appears as a
      `SharedLocalStateRequest`.
    state_requests: A dictionary mapping state names to the corresponding
      `InitialLocalStateRequest` or `FrozenLocalStateRequest`.

  Returns:
    A tree where the first appearance of each shared state is replaced by the
    corresponding `InitialLocalStateRequest` or `FrozenLocalStateRequest`, and
    all other appearances are kept as `SharedLocalStateRequest`.
  
penzai/deprecated/v1/data_effects/local_state.py