Working with pytrees — JAX documentation

JAX has built-in support for objects that look like dictionaries (dicts) of arrays, or lists of lists of dicts, or other nested structures — in JAX these are called pytrees. This section will explain how to use them, provide useful code examples, and point out common “gotchas” and patterns. A pytree is a container-like structure built out of container-like Python objects — “leaf” pytrees and/or more pytrees. A pytree can include lists, tuples, and dicts. A leaf is anything that’s not a pytree, such as an array, but a single leaf is also a pytree. In the context of machine learning (ML), a pytree can contain: Model parameters Dataset entries Reinforcement learning agent observations When working with datasets, you can often come across pytrees (such as lists of lists of dicts). Below is an example of a simple pytree. In JAX, you can use jax.tree.leaves(), to extract the flattened leaves from the trees, as demonstrated here: Any tree-like structure built out of container-like Python obje

JAX has built-in support for objects that look like dictionaries (dicts) of arrays, or lists of lists of dicts, or other nested structures — in JAX these are called pytrees. This section will explain how to use them, provide useful code examples, and point out common “gotchas” and patterns. A pytree is a container-like structure built out of container-like Python objects — “leaf” pytrees and/or more pytrees. A pytree can include lists, tuples, and dicts. A leaf is anything that’s not a pytree, such as an array, but a single leaf is also a pytree. In the context of machine learning (ML), a pytr