🚀 Feature

PyTorch should use a performant and stable version of _pytree.

Motivation

torch.utils._pytree (heavily inspired by JAX pytrees) is something I cooked up in a non-performant way when prototyping the initial non-composable version of torch.vmap. Since then, it's found many legitimate use cases in PyTorch internals:

• functorch transforms; eager compilation
• fx
• writing torch_dispatch objects

The last use case means we want to expose torch.utils._pytree to users! (or provide a stable mechanism that users can use).

There is only really one main blocker to exposing torch.utils._pytree: performance.

• tree_map is slow, it can probably be faster.
• Some of the eager compilation work is showing that torch.utils._pytree is slow, especially in the case where the input is a flat list.

Pitch

torch.utils._pytree is heavily inspired by JAX's pytrees. JAX's pytrees are implemented in C++ (with bindings to Python) for performance reasons. We should either:

• delete our pytree and take a dependency on JAX (or figure out if they'd be interested in splitting off pytree into its own library)
• run our own implementation of pytree and improve its performance.

cc @ezyang @gchanan @kadeng @msaroufim @XuehaiPan @zou3519 @albanD @Chillee