A hierarchy arranges components in nested levels. In JAX, a tree is a nested structure of named or unnamed lists that can hold function state, letting a loss function receive several variables together. section 9.1 section 3.5.3
compute_loss <- function(state, x, y) { #1
.....
loss
}
grad_fn <- jax$value_and_grad(compute_loss)
state <- list(a, b, c)
.[loss, grads_of_loss_wrt_state] <- grad_fn(state, x, y) #2
See state grouped as a list and passed as the first argument to a loss function before computing gradients. section 3.5.3
