derivatives wrt to intermediate variables

[xmax1]

is it possible to get the grads of the output wrt to intermediate variables?

import jax.numpy as jnp
from jax import grad

def function(w, x):
z = w @ x
return jnp.sum(jnp.tanh(z))

dL_dw = grad(function)(w, x)
dL_dz = ??

[xmax1]

Nevermind, it is answered here #5336