def bind_variables(
    tree: Any,
    variables: Iterable[AbstractVariable | AbstractVariableValue],
    allow_unused: bool = False,
    unfreeze_as_copy: bool = False,
) -> Any:
  """Binds variables (mutable or frozen) into the variable slots in a pytree.

  This function re-inserts variable instances into a tree in place of their
  matching variable slots. It is the inverse of `unbind_variables`.

  Args:
    tree: The tree to substitute variables into.
    variables: The collection of variables to insert.
    allow_unused: Whether to ignore variables that do not have any matching slot
      (in which case they will not be inserted).
    unfreeze_as_copy: Whether to unfreeze variable values before inserting them,
      producing a new mutable copy of each input variable. If True, all input
      variables must be instances of `AbstractVariableValue`.

  Returns:
    A copy of ``tree`` with variables re-inserted.
  """
  leaves, treedef = jax.tree_util.tree_flatten(
      tree, is_leaf=lambda l: isinstance(l, AbstractVariableSlot)
  )

  if unfreeze_as_copy:
    orig_variables = variables
    variables = []
    for var in orig_variables:
      if isinstance(var, AbstractVariableValue):
        variables.append(var.unfreeze_as_copy())
      else:
        raise ValueError(
            "unfreeze_as_copy=True is only allowed if all variables are"
            " variable values (e.g. ParameterValue or StateVariableValue)."
        )

  substitution = {}
  for var in variables:
    var_slot = var.get_slot()
    if var_slot in substitution:
      raise ValueError(
          "Cannot substitute multiple variables into the same slot! "
          f"Repeated slot: {var_slot}"
      )
    substitution[var_slot] = var

  seen = set()

  new_leaves = []
  for leaf in leaves:
    if isinstance(leaf, AbstractVariableSlot) and leaf in substitution:
      new_leaves.append(substitution[leaf])
      seen.add(leaf)
    else:
      new_leaves.append(leaf)

  if not allow_unused:
    unseen = set(substitution.keys()) - seen
    if unseen:
      raise ValueError(
          f"Did not find any matching slots for some variables: {unseen}"
      )

  return treedef.unflatten(new_leaves)
Binds variables (mutable or frozen) into the variable slots in a pytree.

  This function re-inserts variable instances into a tree in place of their
  matching variable slots. It is the inverse of `unbind_variables`.

  Args:
    tree: The tree to substitute variables into.
    variables: The collection of variables to insert.
    allow_unused: Whether to ignore variables that do not have any matching slot
      (in which case they will not be inserted).
    unfreeze_as_copy: Whether to unfreeze variable values before inserting them,
      producing a new mutable copy of each input variable. If True, all input
      variables must be instances of `AbstractVariableValue`.

  Returns:
    A copy of ``tree`` with variables re-inserted.
  
penzai/core/variables.py