No method for `prepare_gradient` when the first argument is `Constant` or `Cache`

[albertomercurio]

prepare_gradient doesn't work when the first argument of f is either Constant or Cache.

function my_f(a, x)
  return a * sum(abs2, x)
end

function my_f2(x, a)
  return a * sum(abs2, x)
end

a = 0.1
x = rand(100)

backend = AutoMooncake(; config=nothing)

prep = prepare_gradient(my_f, backend, Constant(a), x)  # doesn't work

prep2 = prepare_gradient(my_f2, backend, x, Constant(a)) # works

[yolhan83]

Yes for now, the differentiated variable must go first and only one is supported, this was the only way to mix backends I think ? Maybe Mooncake.jl could have an autodiff kind of thing outside of DI itself @gdalle may know better if its even possible.

[gdalle]

Hi, thanks for opening the issue.

This is not a bug, this is the official API behavior, as explained e.g. here: the differentiated argument first, and then the rest. This limitation is inherited from legacy or forward-mode autodiff backends, which usually support only one argument (ForwardDiff, ReverseDiff, FiniteDiff). It should also be visible in the docstrings (that's what the distinction between x and contexts... means), but if you think we can make it clearer, I'm open to suggestions!

Multiple arguments or other orders might be supported one day, but to do it right I would need a breaking release and a week's worth of work, so don't hold your breath ^^

[gdalle]

I'm gonna go ahead and close this as unplanned since the API will likely stay that way for the foreseeable future, but don't hesitate to reopen if your question hasn't been properly answered!