# Multi-argument Jacobian and gradient
**URL:** https://discourse.julialang.org/t/multi-argument-jacobian-and-gradient/129029
**Category:** Numerics
**Tags:** question, ad
**Created:** [May 15, 2025, 11:06am UTC](https://discourse.julialang.org/t/multi-argument-jacobian-and-gradient/129029 "2025-05-15T11:06:31Z")
**Posts on this page:** 8
**Page:** 1
### Author:  [@Tamas\_Papp](https://discourse.julialang.org/u/Tamas_Papp)
#### Post date: [May 15, 2025, 11:06am UTC](https://discourse.julialang.org/t/multi-argument-jacobian-and-gradient/129029/1 "2025-05-15T11:06:31Z")
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?
```julia
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)
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [May 15, 2025, 11:10am UTC](https://discourse.julialang.org/t/multi-argument-jacobian-and-gradient/129029/2 "2025-05-15T11:10:46Z")
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](https://github.com/EnzymeAD/Enzyme.jl/issues/1950)). A less manual alternative would be ComponentArrays.jl.
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [May 15, 2025, 11:18am UTC](https://discourse.julialang.org/t/multi-argument-jacobian-and-gradient/129029/3 "2025-05-15T11:18:45Z")
On second thought the views problem shouldn’t affect you there, so maybe try both and benchmark
---
### Author:  [@Tamas\_Papp](https://discourse.julialang.org/u/Tamas_Papp)
#### Post date: [May 15, 2025, 12:56pm UTC](https://discourse.julialang.org/t/multi-argument-jacobian-and-gradient/129029/4 "2025-05-15T12:56:47Z")
Incidentally, the manual suggests problems with closures. Are they still valid? Does using `Base.Fix1` etc help there?
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [May 15, 2025, 2:56pm UTC](https://discourse.julialang.org/t/multi-argument-jacobian-and-gradient/129029/5 "2025-05-15T14:56:52Z")
Assuming you don’t run into the infamous [captured variable](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-captured) 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.
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [May 15, 2025, 2:57pm UTC](https://discourse.julialang.org/t/multi-argument-jacobian-and-gradient/129029/6 "2025-05-15T14:57:43Z")
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](https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/api/#DifferentiationInterface.FixTail).
---
### Author:  [@jonniedie](https://discourse.julialang.org/u/jonniedie)
#### Post date: [May 15, 2025, 11:27pm UTC](https://discourse.julialang.org/t/multi-argument-jacobian-and-gradient/129029/7 "2025-05-15T23:27:40Z")
For index bookkeeping, `ComponentArrays` can be pretty helpful as it also lets you also access blocks of your Jacobian by name (note, the Jacobian only seems to retain its `ComponentArray`-ness for in-place calculations, as it gets allocated as a plain `Matrix` otherwise):
```julia
using ComponentArrays
import ForwardDiff
import DifferentiationInterface as DI
const A = randn(3, 3)
function f!(output, input)
(; x, y) = input
# Our output has two named components
output.a = x' * y
output.b .= A*x .+ A'*y
end
# Define our inputs as having two named 3-vector components
input = ComponentArray(x=ones(3), y=ones(3))
# Preallocate an output for our in-place calculations
output = ComponentArray(a=0.0, b=zeros(3))
# Preallocate a Jacobian, which will be a ComponentMatrix with
# named row indices a=1, b=2:4
# named column indices x=1:3, y=4:6
J = 0 * output .* input'
# Calculate the Jacobian
DI.jacobian!(f!, output, J, DI.AutoForwardDiff(), input)
# We can access blocks of the Jacobian by name
J[Val(:b), Val(:y)]
# prints
# 3×3 Matrix{Float64}:
# -0.671587 0.0719139 -0.0111287
# 0.972351 -1.07426 -1.03
# 0.718292 1.59966 0.148113
J[Val(:a), Val(:x)]
# prints
# 3-element Vector{Float64}:
# 1.0
# 1.0
# 1.0
```
---
### Author:  [@Tamas\_Papp](https://discourse.julialang.org/u/Tamas_Papp)
#### Post date: [May 21, 2025, 12:04pm UTC](https://discourse.julialang.org/t/multi-argument-jacobian-and-gradient/129029/8 "2025-05-21T12:04:31Z")
> [@jonniedie](#):
>
> note, the Jacobian only seems to retain its `ComponentArray`-ness for in-place calculations
Thanks for this, it is a key part of what I was missing. (I tried it with the non-! methods and always got a plain vanilla `Array`).