Taking derivatives with respect to only some elements of an array

[raphaelchinchilla]

This is a feature suggestion. I do not believe it currently exists.

Suppose one has access to a function f(x,p) that, and this is the important part, cannot be rewritten or modified. In my particular example, I am solving a differential equation, but another situation could be:

using Flux, Optimisers
model=Dense(2,2)
p,re=destructure(model)
function f(x,p)
  m = re(p)
  return m(x)
end

For simplicity, suppose p is a vector and let us split it into p=[p1;p2]. I would like to take the gradient of a cost function of f(x,p) (let us say, sum) with respect to only p1. I think that the best way to do it currently is something like

gradient(p1) do p1
  p=[p1;p2]
  sum(f(x,p))
end

However, this allocates a new vector p each time which is inefficient, specially if the number of elements in p1 is much smaller than the one in p2 (another possibility is to take the gradient with respect to both p1 and p2 and ignore the p2 component but this is probably even more inefficient).

For context, I am needing to take the derivative with respect to only some parameters because I am updating the others with another method (that does not require the gradient).

I am not sure how, but it would be nice if there were a way to tell Zygote to only compute the gradient with respect to some indexes of a vector (and also the convert function, to not compute the gradient with respect to some indexes)

[ToucheSir]

#981 is one thing which would help this, but it is unlikely to be addressed holistically because Zygote's analysis is just not that smart (the exact opposite, really). In fact, it does not even know at "compile time" whether certain values are arrays or not (hence why #981 is purely a "runtime" solution).

[raphaelchinchilla]

Oh, what a shame.

Should I keep the issue open or just go ahead and close it?

[ToucheSir]

We can leave it open for now and see if your use case is addressed by that PR if/when it's merged.