utility apply reduce function to a pytree (or a pytree of pytree)

tree_reduce.py

import numpy as np

from jax import numpy as jnp
from jax.lib import pytree


def reduce(trees, op=jnp.stack, op_kwargs: dict = {}):
    """Takes a list of trees and apply op to every corresponding leaf.
    For example, assuming op=`jnp.stack`,
    given two trees ((a, b), c) and ((a', b'), c'), returns
    ((stack(a, a'), stack(b, b')), stack(c, c')).
    # Start as a fork of: https://gist.github.com/willwhitney/dd89cac6a5b771ccff18b06b33372c75
    """
    leaves_list = []
    treedef_list = []
    for tree in trees:
        leaves, treedef = pytree.flatten(tree)
        leaves_list.append(leaves)
        treedef_list.append(treedef)

    grouped_leaves = zip(*leaves_list)
    result_leaves = [op(l, **op_kwargs) for l in grouped_leaves]
    return treedef_list[0].unflatten(result_leaves)