def combine(*partitions: Any) -> Any:
  """Combines leaves from multiple partitions.

  This function can be used to reverse the results of calling `partition` on
  a selector. It walks two or more PyTrees, and combines their leaves by
  detecting and replacing instances of `NotInThisPartition`. The partitions can
  have `NotInThisPartition` replacing any subtree, not just leaves.

  A common use case for this function is to recombine parts of an input model
  after splitting them to handle them differently in JAX transformations. See
  the documentation for `penzai.core.selectors.Selection.partition` for more
  details.

  This function is inspired by Equinox's `eqx.combine`, but uses
  `NotInThisPartition` as the sentinel instead of None, and supports
  `NotInThisPartition` at arbitrary subtree locations instead of only having
  None at leaves. (Partitioning is also somewhat less important in Penzai than
  in Equinox because all PyTree leaves are arraylike by convention; partitioning
  is only necessary when different parts of the tree need special treatment.)

  Args:
    *partitions: Partitions to combine. All partitions should have the same
      PyTree structure except that, for each PyTree leaf, exactly one of the
      input partitions actually has that leaf, and every other partition has
      `NotInThisPartition` sentinels in place of that leaf or one of its
      ancestors.

  Returns:
    Combined version of all partitions, which takes the concrete value from
    each partition instead of the `NotInThisPartition` sentinel.
  """
  if not partitions:
    # NotInThisPartition() is the identity for `combine` and thus a reasonable
    # return value for combining the empty set.
    return NotInThisPartition()

  def combine_pairwise_subtrees(keypath, tree_a, tree_b):
    if isinstance(tree_a, NotInThisPartition):
      return tree_b
    elif isinstance(tree_b, NotInThisPartition):
      return tree_a
    else:
      flattened_a = tree_util.tree_flatten_exactly_one_level(tree_a)
      flattened_b = tree_util.tree_flatten_exactly_one_level(tree_b)
      if flattened_a is not None and flattened_b is not None:
        keys_and_subtrees_a, treedef_a = flattened_a
        keys_and_subtrees_b, treedef_b = flattened_b
        if treedef_a == treedef_b:
          combined_children = []
          for (key_a, subtree_a), (key_b, subtree_b) in zip(
              keys_and_subtrees_a, keys_and_subtrees_b
          ):
            assert key_a == key_b
            combined_children.append(
                combine_pairwise_subtrees(
                    keypath + (key_a,), subtree_a, subtree_b
                )
            )
          return treedef_a.unflatten(combined_children)

      # If we got here, tree structures aren't compatible.
      raise ValueError(
          f"Expected the nodes at path {jax.tree_util.keystr(keypath)} to have"
          " compatible structures, or for one of them to be"
          f" NotInThisPartition. However, got {tree_a} and {tree_b}, which are"
          " incompatible."
      )

  return functools.reduce(
      functools.partial(combine_pairwise_subtrees, ()), partitions
  )
Combines leaves from multiple partitions.

  This function can be used to reverse the results of calling `partition` on
  a selector. It walks two or more PyTrees, and combines their leaves by
  detecting and replacing instances of `NotInThisPartition`. The partitions can
  have `NotInThisPartition` replacing any subtree, not just leaves.

  A common use case for this function is to recombine parts of an input model
  after splitting them to handle them differently in JAX transformations. See
  the documentation for `penzai.core.selectors.Selection.partition` for more
  details.

  This function is inspired by Equinox's `eqx.combine`, but uses
  `NotInThisPartition` as the sentinel instead of None, and supports
  `NotInThisPartition` at arbitrary subtree locations instead of only having
  None at leaves. (Partitioning is also somewhat less important in Penzai than
  in Equinox because all PyTree leaves are arraylike by convention; partitioning
  is only necessary when different parts of the tree need special treatment.)

  Args:
    *partitions: Partitions to combine. All partitions should have the same
      PyTree structure except that, for each PyTree leaf, exactly one of the
      input partitions actually has that leaf, and every other partition has
      `NotInThisPartition` sentinels in place of that leaf or one of its
      ancestors.

  Returns:
    Combined version of all partitions, which takes the concrete value from
    each partition instead of the `NotInThisPartition` sentinel.
  
penzai/core/partitioning.py