Suppose we have a function f(x, y), and want to compute \partial_x f(x,y) for many y’s, ideally using no memory allocation. Using Zygote, I tried Zygote.forwarddiff(y->f(x,y), y).

However, I believe this creates a new anonymous function each time we call it with a different x. So the call to fowarddiff is not efficient.

Also, we can’t use gradient or forwarddiff on f directly, because in reality there are many more arguments (x,y,z,…) and computing and storing all of their derivatives is again inefficient.

Is there a way to achieve what I intend to do with Zygote or any other AD package?