Method partition

penzai/core/selectors.py:964–1032

Partitions the tree into ``(selected_tree, remainder_tree)`` parts. This function can be used to separate out the selected components of a tree into their own separate tree, so that JAX functions and other JAX libraries can process them like ordinary PyTrees. It splits its input into tw

(self, at_leaves: bool = False)

    return self.at_subtrees_where(_check_equal)

  def partition(self, at_leaves: bool = False) -> tuple[Any, Any]:
    """Partitions the tree into ``(selected_tree, remainder_tree)`` parts.

    This function can be used to separate out the selected components of a tree
    into their own separate tree, so that JAX functions and other JAX libraries
    can process them like ordinary PyTrees. It splits its input into two
    disjoint trees ``(selected_tree, remainder_tree)``, where ``selected_tree``
    only contains the leaves that were selected, and `remainder_tree` only
    contains the remainder. The parts that were removed are identified using a
    sentinel `pz.NotInThisPartition` object, which has no PyTree children.

    The main use case for ``partition`` is to identify subsets of models that
    should be treated in different ways by JAX API functions. For instance, if
    you want to take a gradient with respect to a specific subset of parameters,
    you can select those parameters, call ``partition`` to separate them from
    the rest, then call `jax.grad` and use `argnums` to identify the partition
    of interest. Similarly, if you want to donate only a subset of the state to
    `jax.jit`, you can partition it and then use JAX's ``donate_argnums``
    argument to `jax.jit` to identify the parts you want to donate. Inside the
    function, you can then use `pz.combine` to rebuild the original tree.

    It is possible to repeatedly call ``partition`` to split a tree into
    more than two parts. In particular, you can select the ``remainder_tree``,
    target some additional nodes, and call ``.partition()`` again, repeating
    this process as needed. All of the partitioned trees can then be re-combined
    using a single call to `pz.combine`.

    Note that `NotInThisPartition` is a PyTree node with no children, which
    means that partitioned trees are safe to pass through JAX transformations,
    and the set of leaves in the two partitioned trees together are the same as
    the set of leaves in the original selected tree.

    This function is inspired by Equinox's `equinox.partition`, but is designed
    to work with Penzai's selector system. Unlike `equinox.partition`, missing
    nodes are identified with the `pz.NotInThisPartition` sentinel, and can
    replace arbitrary PyTree subtrees instead of just 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 e.g. for ``argnums`` or
    ``donate_argnums`` parameters.)

    Args:
      at_leaves: Whether to do the partitioning at the leaf level, so that the
        returned trees have exactly the same structure. (Note that `pz.combine`
        is OK with entire subtrees missing, so this is not necessary, but can
        make the partitions easier to manipulate manually if desired.) If False,
        the entire selected subtrees will be replaced by `NotInThisPartition` in
        the remainder tree.

    Returns:
      A tuple ``(selected_tree, remainder_tree)``, where both trees have the
      same structure (if ``at_leaves=True``) or the same prefix (if
      ``at_leaves=False``) except that `NotInThisPartition` is used to replace
      parts that are in the other partition.
    """
    if at_leaves:
      selected_tree = (
          self.invert()