3 / 8 JAX & XLA Mental Model

Tree Hugging

A model's parameters are stored as a nested dict of dicts holding arrays, e.g. {"layer1": {"w": W1, "b": b1}, "layer2": {"w": W2, "b": b2}}. What JAX concept lets transforms like jax.grad operate on this whole nested structure as if it were a single array?

A

A custom jax.Array subclass that overrides add

B

Pytree, JAX's generic abstraction for nested Python containers whose leaves are arrays

C

A NamedTuple schema registered via jax.jit's static_argnums

D

jax.numpy.stack, which flattens nested dicts into one array before tracing