VJPs of VJPs?

[karan-dalal]

Reducing the scope of implementation to make the problem more understandable.

I have a function that takes in an input x, model parameters W, and produces and output y. Loss is computed as the MSE between x and y. My goal is to take the derivative of the "gradient of loss WRT the model" wrt to the output. To illustrate mathematically:
$$\frac{d}{dy}(\frac{d loss}{dW})$$

My thought process for implementation is to decompose the $\frac{d loss}{dy}$ and $\frac{dy}{d W}$ into a VJP with upstream gradient, then take the derivative of this:
$$\frac{d loss}{dW} = \frac{d loss}{d y} * \frac{d y}{d W}$$

Implemented as:

def f(x, W):
    y = x @ W
    return y

def loss(x, y):
    return MSE(x, y)

def fn(x, W):
    loss_fn = grad(loss, argnums=1)
    y, f_vjp = vjp(f, x, W)
    
    def nest_vjp(y):
        d_loss_dy = loss_fn(x, y)
        d_loss_dW, _ = f_vjp(d_loss_dy)
        return d_loss_dW

    d_loss_dW, nested_vjp_fn = vjp(nest_vjp, y)
    d_L_dy = nested_vjp_fn(# Some upstream gradient)

The goal for this is to compute some outer loop (upstream) gradient with respect to the output of an inner loop. Is it possible to nest VJPs as such – and will it compute the correct $\frac{d}{dy}(\frac{d loss}{dW})$?

Thanks, let me know if I can clarify anything!

[karan-dalal]

@jakevdp @patrick-kidger any help appreciated!

[karan-dalal]

Just wanted to follow up here @jakevdp Thank you.

[jakevdp]

It seems like #18383 is a better way to pose this question.