def hoist_constant_side_inputs(
    tree: Any,
    hoist_predicate: Callable[[Tag, Any], bool] = lambda tag, val: True,
    scoped: bool = True,
) -> tuple[Any, dict[Tag, Any]]:
  """Extracts all constant side inputs from a tree so they can be re-handled.

  This function finds all `WithConstantSideInputs` handlers in a tree, and
  extracts the side inputs from those handlers, and replaces the bound side
  input references with unhandled requests. You can then manipulate the tree
  and finally re-handle the side inputs using `WithConstantSideInputs.handling`.

  This function is useful when a tree contains constant side inputs inside a
  subtree, and you want to pull out an even smaller subtree that uses those
  side inputs, or move around subtrees that use side inputs. A particular
  instance of this is shared parameters, which can be represented in Penzai as
  side inputs with a constant value provided by a `WithConstantSideInputs`
  handler. If you want to extract a layer that uses shared parameters, you can
  first hoist out the constant values of all shared parameters, extract the
  submodel you want, and then re-bind the shared parameters of that subtree
  alone.

  Calling `hoist_constant_side_inputs` on a layer and then immediately calling
  `WithConstantSideInputs.handling` on the results will usually produce a layer
  with the same observable behavior, although the tree structure may differ.
  However, if ``scoped`` is False, it is possible that this function will raise
  an error due to tag conflicts.

  Args:
    tree: A tree to hoist constant side inputs from. Usually will contain
      WithConstantSideInputs handlers somewhere inside it.
    hoist_predicate: Optional predicate to determine which side inputs to try to
      hoist, given the tag and value. Defaults to hoisting all side inputs.
    scoped: Whether to modify the tags for the hoisted side inputs so that they
      remain unique even if different side inputs use different tags.

  Returns:
    Tuple `(effectful_tree, hoisted_values)` where `effectful_tree` is a copy
    of `tree` with all hoisted side inputs replaced by `SideInputRequest`
    instances and all `WithConstantSideInputs` removed or modified to not
    include the hoisted side inputs, and `hoisted_values` are the values that
    were bound by those constant side inputs, as needed by
    `WithConstantSideInputs.handling`.
  """
  hoisted_values = {}
  handler_remaps = {}

  def _prepare(subtree, keypath_prefix):
    for keypath, handler in (
        selectors.select(subtree)
        .at_instances_of(WithConstantSideInputs)
        .get_by_path()
        .items()
    ):
      if scoped:
        source = tree_util.pretty_keystr(keypath_prefix + keypath, tree)
        wrap_tag = functools.partial(HoistedTag, source)
      else:
        wrap_tag = lambda tag: tag

      cur_remap = {}
      for tag, value in handler.side_inputs.items():
        if hoist_predicate(tag, value):
          new_tag = wrap_tag(tag)
          if new_tag in hoisted_values:
            raise ValueError(
                f"Tag conflict while hoisting tag {repr(new_tag)}! Try using"
                " `scoped` to ensure the tags are unique."
            )
          hoisted_values[new_tag] = value
          cur_remap[tag] = new_tag
      handler_remaps[handler.handler_id] = cur_remap

      _prepare(
          handler.body,
          keypath_prefix + keypath + (jax.tree_util.GetAttrKey("body"),),
      )

  _prepare(tree, ())

  def _find_and_hoist(subtree):
    return (
        selectors.select(subtree)
        .at_instances_of(
            SideInputRequest | HandledSideInputRef | WithConstantSideInputs
        )
        .apply(_hoister)
    )

  def _hoister(
      node: SideInputRequest | HandledSideInputRef | WithConstantSideInputs,
  ):
    if isinstance(node, SideInputRequest):
      if node.tag in hoisted_values:
        raise ValueError(
            "Cannot hoist a tag that is already used by an unhandled"
            f" SideInputRequest! Found an unhandled request with tag {node.tag}"
        )
      else:
        return node
    elif isinstance(node, HandledSideInputRef):
      if (
          node.handler_id in handler_remaps
          and node.tag in handler_remaps[node.handler_id]
      ):
        return SideInputRequest(handler_remaps[node.handler_id][node.tag])
      else:
        return node
    elif isinstance(node, WithConstantSideInputs):
      assert node.handler_id in handler_remaps
      cur_remap = handler_remaps[node.handler_id]
      # Process the body recursively.
      processed_body = _find_and_hoist(node.body)
      # Hoist any relevant side inputs.
      new_side_inputs = {}
      for tag, value in node.side_inputs.items():
        if tag not in cur_remap:
          new_side_inputs[tag] = value
      if new_side_inputs:
        return WithConstantSideInputs(
            handler_id=node.handler_id,
            body=processed_body,
            side_inputs=new_side_inputs,
        )
      else:
        # Every tag this handler provides has been hoisted.
        return processed_body
    else:
      raise TypeError(f"unexpected node {node}")

  final_tree = _find_and_hoist(tree)
  return final_tree, hoisted_values
Extracts all constant side inputs from a tree so they can be re-handled.

  This function finds all `WithConstantSideInputs` handlers in a tree, and
  extracts the side inputs from those handlers, and replaces the bound side
  input references with unhandled requests. You can then manipulate the tree
  and finally re-handle the side inputs using `WithConstantSideInputs.handling`.

  This function is useful when a tree contains constant side inputs inside a
  subtree, and you want to pull out an even smaller subtree that uses those
  side inputs, or move around subtrees that use side inputs. A particular
  instance of this is shared parameters, which can be represented in Penzai as
  side inputs with a constant value provided by a `WithConstantSideInputs`
  handler. If you want to extract a layer that uses shared parameters, you can
  first hoist out the constant values of all shared parameters, extract the
  submodel you want, and then re-bind the shared parameters of that subtree
  alone.

  Calling `hoist_constant_side_inputs` on a layer and then immediately calling
  `WithConstantSideInputs.handling` on the results will usually produce a layer
  with the same observable behavior, although the tree structure may differ.
  However, if ``scoped`` is False, it is possible that this function will raise
  an error due to tag conflicts.

  Args:
    tree: A tree to hoist constant side inputs from. Usually will contain
      WithConstantSideInputs handlers somewhere inside it.
    hoist_predicate: Optional predicate to determine which side inputs to try to
      hoist, given the tag and value. Defaults to hoisting all side inputs.
    scoped: Whether to modify the tags for the hoisted side inputs so that they
      remain unique even if different side inputs use different tags.

  Returns:
    Tuple `(effectful_tree, hoisted_values)` where `effectful_tree` is a copy
    of `tree` with all hoisted side inputs replaced by `SideInputRequest`
    instances and all `WithConstantSideInputs` removed or modified to not
    include the hoisted side inputs, and `hoisted_values` are the values that
    were bound by those constant side inputs, as needed by
    `WithConstantSideInputs.handling`.
  
penzai/deprecated/v1/data_effects/side_input.py