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?
Sign in to answer questions and track your progress
Sign In