embed_shared_state_requests

penzai.data_effects.local_state.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.

Parameters:

• 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.