liveBook · Manning

Hierarchy

As explained in Deep Learning with Python, Third Edition

A nested collection structure of tuples, dicts, and lists used by JAX to represent composite inputs and states.

In JAX, a tree is a state organized as a tuple, dictionary, list, or nested combination of these. It lets a loss function accept several state variables together while keeping other inputs separate. section 3.5.3

def compute_loss(state, x, y):                                     #1
    ...
    return loss

grad_fn = jax.value_and_grad(compute_loss)
state = (a, b, c)
loss, grads_of_loss_wrt_state = grad_fn(state, x, y)                #2

The example passes a three-variable tuple as one state argument to a loss function, then obtains gradients for that state. section 3.5.3

How Deep Learning with Python, Third Edition covers it

How deep it goes
Defines trees as nested tuples, dictionaries, or lists and demonstrates their use as function state.
Examples
Uses a tuple containing a, b, and c as state, with x and y as separate loss inputs.
What you can do after
Apply it by passing a tuple, dictionary, list, or nested combination as state to a JAX loss-gradient function.
How it connects
Links JAX trees to loss functions, state arguments, and gradient computation with jax.value_and_grad().
Cover of Deep Learning with Python, Third Edition
liveBookmanning.com
Read Deep Learning with Python, Third Edition on liveBook Get the book at manning.com

Also explained in

Compare how 7 books explain hierarchy →

Related concepts

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with Python, Third Edition. The section links point to the parts of the book it draws on.