When model parameters, dataset elements, or results become nested combinations of lists, dictionaries, tuples, and arrays, passing every value separately makes code harder to read. A pytree packages that hierarchy so JAX functions and transformations can process it as one value. section 10.0 section 2.11
A pytree is a recursive structure of nodes and leaves. Registered containers such as list, tuple, dict, namedtuple, OrderedDict, and None form nodes; numeric values, JAX arrays, and NumPy arrays are leaves by default. Other container classes can be registered with flattening and unflattening rules. section 10.1 section 10.4
Use jax.tree_util to inspect, transform, and reorganize nested values:
- tree_map() applies a function to corresponding leaves and returns a new pytree; it does not modify the original in place. section 10.2.1
- tree_leaves() extracts leaves, while tree_structure() produces a PyTreeDef describing the hierarchy. section 10.4
- tree_flatten() and tree_unflatten() separate leaves from structure and reconstruct the original pytree; ravel_pytree() creates a flat one-dimensional representation with an unflattening function. section 10.4
- tree_reduce() combines leaves into one value, and tree_transpose() changes an outer-then-inner hierarchy into an inner-then-outer hierarchy. section 10.2.3 section 10.2.4
Mapping over neural-network parameters
A multilayer network can store each layer as a tuple of weight matrix and bias vector, with all layers held in a list. Applying `jax.tree_util.tree_map(lambda p: p.shape, params)` produces the same hierarchy containing shapes instead of arrays. Replacing the function with `lambda p: 10 * p` creates scaled parameters. The same pattern can broadcast every parameter across a device axis before `pmap()` training. section 10.1 section 10.2.1
Preparing structured data for vmap
A list of namedtuple points is an array of structures, which `vmap()` cannot map efficiently along axis 0. Use `tree_transpose()` to convert it into a structure of arrays, such as `Point(x=[...], y=[...])`, then convert the coordinate lists to JAX arrays. The resulting structure can be passed to `vmap()` for rotating every point. section 10.2.4
