pytree transformation

I’d like to write code like np.sum((x - y) ** 2) that automatically works on arbitrarily nested Python objects.

I think this could be done cleanly with a “pytree transformation” that acts on functions, somewhat similar to the existing vmap, e.g.,

@jax.autotree
def norm(x, y):
    return np.sum((x - y) ** 2)

Under the hood, we would need to define what “autotree” versions of each primitive function look like on a list of heterogeneously shaped arrays, e.g.,

• Vectorized functions like + operate over each array independently.
• In mixed tree/array operations, the array should get “broadcast” by repeating it when the structure of the tree, e.g., {'a': x, 'b': y'} + z -> {'a': x + z, 'b': y + z}.
• Reductions like sum sum over all the arrays to produce a scalar, unless an axis is specified.

The upside is that this could make it trivial to write functional transformations like jax.experimental.optimizers in a fully general way. I can think of lots of nice applications for this, e.g., for methods that solve equations.

[shoyer]

I need to figure out how to get #3263 moving towards something we can merge. The main unresolved challenge to make this useful is figuring out how to get JAX transformations and control flow inside tree vectorized functions to work. There is a lot of reorganizing data for new calling conventions that gets rather tedious, but maybe we can just power through it or figure out a better way, e.g., perhaps with partial tree flattening using the new(ish) is_leaf parameter.

With regard to in_trees/out_trees, that seems like a very sensible suggestion (similar to in_axes/out_axes for vmap), although you can get a similar effect to in_trees just by defining an inner function to apply tree_vectorize to using a closure. I think this is somewhat more of a superficial issue, though – the key challenge to tackle first is figuring out how to get the fundamental JAX transformations inside tree_vectorize working.

CC @apaszke

[shoyer]

Update: I have a new approach for implementing this, which I think is much more promising than #3263:

• #8504

Feedback would be very welcome!