Multi-argument Jacobian and gradient

Specific Domains Numerics
,
1

I have a function that takes multiple arguments (all of them vectors), and I want to obtain the Jacobian in each parameter (or, similarly, the gradient if the function maps to \mathbb{R}).

Is it best to unpack/pack like this, or is there a better way?

import ForwardDiff
import DifferentiationInterface as DI

const A = randn(3, 3)
f(x, y) = A*x .+ A'*y
J = DI.jacobian(v -> f(@view(v[1:3]), @view(v[4:6])), DI.AutoForwardDiff(), ones(6))
J[:, 1:3], J[:, 4:6]

Is there a package to help with the index bookkeeping? (Which is trivial in this simple case, but would be better to do it right if I have more arguments)

2

First of all, some AD packages can get you the gradient/Jacobian with respect to each argument natively (like Zygote, or Enzyme). in DI, we chose to only support one active argument, which means you indeed have to pack somehow if you want to leverage the interface.
The manual way you chose is not bad, but I think you might run into some trouble with the views (see e.g. this Enzyme issue). A less manual alternative would be ComponentArrays.jl.

1 Like
3

On second thought the views problem shouldn’t affect you there, so maybe try both and benchmark

1 Like
4

Incidentally, the manual suggests problems with closures. Are they still valid? Does using Base.Fix1 etc help there?

5

Assuming you don’t run into the infamous captured variable issue, I think Enzyme.jl is the only backend for which the presence of closures might impact performance. Even so, I don’t know how big the impact would be in real life, and it probably depends on the function.
Callable structs like Base.Fix1 can help you avoid the captured variable issue, but otherwise they don’t fundamentally behave differently.

1 Like
6

In fact, the handling of context arguments with most backends is not available out of the box, so DI implements it with Base.Fix1 or similar tricks.

1 Like