PyTrees Are Not One Thing: JAX, PyTorch, and TensorFlow Compared

PyTrees look deceptively simple. You flatten a nested Python object into leaves, keep a structure descriptor, and later rebuild or map over the same shape. That abstraction is powerful enough to carry optimizer states, model parameters, batched inputs, gradients, and sharding annotations. It is also just ambiguous enough that three major frameworks implement three subtly different languages under the same idea. This note compares JAX jax.tree_util , PyTorch torch.utils._pytree , and TensorFlow tf.nest . I tested the behavior in two environments: an older stack with JAX 0.4.35, PyTorch 2.2.2, TensorFlow 2.20.0, and a newer stack with JAX 0.10.0, PyTorch 2.12.0, TensorFlow 2.21.0. Most flatten/unflatten semantics were stable across these versions. The main version-sensitive result is PyTorch: _pytree.tree_map in 2.2.2 accepts only one pytree, while 2.12.0 supports multiple pytrees and behaves much closer to JAX prefix-style mapping. The short version: JAX treats pytrees as a transformation language, PyTorch is converging toward that model in torch.func , and TensorFlow exposes a broader nested-structure utility through tf.nest . Those differences show up exactly where backend-agnostic libraries usually hurt: None , dictionary order, custom containers, tree_map , autodiff, and vectorization. The Shape Of The APIs The three APIs have the same surface story but not the same contract. from jax import tree_util as jtu from torch.utils import _pytree as tpu import tensorflow as tf leaves , treedef = jtu . tree_flatten ( tree ) tree = jtu . tree_unflatten ( treedef , leaves ) tree = jtu . tree_map ( f , * trees ) leaves , spec = tpu . tree_flatten ( tree ) tree = tpu . tree_unflatten ( leaves , spec ) tree = tpu . tree_map ( f , tree ) # PyTorch 2.2.2 tree = tpu . tree_map ( f , * trees ) # PyTorch 2.12.0 leaves = tf . nest . flatten ( tree ) tree = tf . nest . pack_sequence_as ( structure , leaves ) tree = tf . nest . map_structure ( f , * structures ) Flattening means "whi