# Preserving ComponentArray structure with DifferentiationInterface

**URL:** <https://discourse.julialang.org/t/preserving-componentarray-structure-with-differentiationinterface/133390>\
**Category:** General Usage\
**Tags:** question, autodiff, dsge\
**Created:** [October 23, 2025, 3:21pm UTC](https://discourse.julialang.org/t/preserving-componentarray-structure-with-differentiationinterface/133390 "2025-10-23T15:21:51Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![charlesknipp](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/charlesknipp/32/214437_2.png) [@charlesknipp](https://discourse.julialang.org/u/charlesknipp)\
**Post date:** [October 23, 2025, 3:21pm UTC](https://discourse.julialang.org/t/preserving-componentarray-structure-with-differentiationinterface/133390/1 "2025-10-23T15:21:51Z")

</div>

I’m building a DSGE modeling language (lol what else is new) where I want to use something like ComponentArrays.jl for the construction of first order perturbations using DifferentiationInterface.jl. The point here is that I divide a model into blocks with unknowns and targets ala [sequence\_jacobian.py](https://github.com/shade-econ/sequence-jacobian).

The problem here is that the component array loses its structure when performing `DI.jacobian()`, but not with `DI.gradient()`. I can produce a (not so) minimum working example with the following.

```julia-auto
using ComponentArrays
using DifferentiationInterface
import ForwardDiff

lag(x::AbstractVector) = x[1]
lag(x::Real) = x

contemp(x::AbstractVector) = x[2]
contemp(x::Real) = x

lead(x::AbstractVector) = x[3]
lead(x::Real) = x

function production(x::ComponentArray)
    (; K, L, Z, α, δ) = x
    r = contemp(α) * contemp(Z) * (lag(K) / contemp(L)) ^ (contemp(α) - 1) - contemp(δ)
    w = (1 - contemp(α)) * contemp(Z) * (lag(K) / contemp(L)) ^ contemp(α)
    Y = contemp(Z) * lag(K) ^ contemp(α) * contemp(L) ^ (contemp(α) - 1)
    return ComponentArray(; r, w, Y)
end

```

Suppose we already solved for the steady state using NonlinearSolve.jl using `x -> production(x)` as the objective, which gave us

```julia-auto
ss = ComponentArray(
    L = 1.0, δ = 0.025, α = 0.11, K = 3.142857142857058, Z = 0.881646097521458
)

```

Finding the first order perturbation then reduces to the following calculus

```julia-auto
stacked_ss = ComponentArray(NamedTuple{keys(ss)}([fill(i, 3) for i in ss]))
J = DifferentiationInterface.jacobian(
    x -> production(x), AutoForwardDiff(), stacked_ss
)

```

which yields the following illegible mess:

```julia-auto
3×15 Matrix{Float64}:
 -0.00991 0.0 0.0 0.0 0.03115 0.0 0.0 0.03969 0.0 0.0 0.35826 0.0 0.0 -1.0 0.0
  0.03115 0.0 0.0 0.0 -0.09790 0.0 0.0 1.00948 0.0 0.0 0.01917 0.0 0.0 0.0 0.0
  0.03500 0.0 0.0 0.0 -0.89000 0.0 0.0 1.13424 0.0 0.0 1.14513 0.0 0.0 0.0 0.0

```

However, when I individually take the gradient per target…

```julia-auto
DifferentiationInterface.gradient(
    x -> production(x)[:r], AutoForwardDiff(), stacked_ss
)

```

the component array structure is preserved!

```julia-auto
ComponentVector{Float64}(
  L = [0.0, 0.031150000000000788, 0.0],
  δ = [0.0, -1.0, 0.0],
  α = [0.0, 0.35826144883243144, 0.0],
  K = [-0.009911363636364154, 0.0, 0.0],
  Z = [0.0, 0.0396984686921376, 0.0]
)

```

As much as I enjoy the type stability and composability of ComponentArrays.jl, I feel like I have to do more work to step away from the intuitiveness of named tuples. And if I have to reshape and manipulate every output from the function call, then I may as well create my own struct specific to this problem.

There are a couple hills I think are worth dying on here:

1. `DI.jacobian` is the clear winner here since we can prep the gradient upon block definition and store sparsity patterns.
2. I need to preserve the names of rows/columns of the Jacobian for computing chain rules when faced with multiple blocks. Whether this be ComponentArrays.jl, NamedTuples, or some custom object.

For those who are curious, I copied the design philosophy of [sequence\_jacobian.py](https://github.com/shade-econ/sequence-jacobian) and create the above model block with a macro.

```julia-auto
@simple function production(K, L, Z, α, δ)
    r = α * Z * (lag(K) / L) ^ (α - 1) - δ
    w = (1 - α) * Z * (lag(K) / L) ^ α
    Y = Z * lag(K) ^ α * L ^ (α - 1)
    return r, w, Y
end

```

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 23, 2025, 4:59pm UTC](https://discourse.julialang.org/t/preserving-componentarray-structure-with-differentiationinterface/133390/2 "2025-10-23T16:59:54Z")

</div>

Hi, glad that you find DI useful! A few things are worth noting:

- The behavior you describe is not due to DI, it comes from the underlying autodiff backend.

```julia
julia> using ComponentArrays, Enzyme, ForwardDiff, Zygote

julia> f(x::ComponentVector) = ComponentVector(c=2 .* x.a, d=3 .* x.b)
f (generic function with 1 method)

julia> x = ComponentVector(a=[1.0, 2.0], b=[3.0, 4.0])
ComponentVector{Float64}(a = [1.0, 2.0], b = [3.0, 4.0])

julia> y = f(x)
ComponentVector{Float64}(c = [2.0, 4.0], d = [9.0, 12.0])

julia> ForwardDiff.jacobian(f, x)
4×4 Matrix{Float64}:
 2.0 0.0 0.0 0.0
 0.0 2.0 0.0 0.0
 0.0 0.0 3.0 0.0
 0.0 0.0 0.0 3.0

julia> Zygote.jacobian(f, x)[1]
4×4 Matrix{Float64}:
 2.0 0.0 0.0 0.0
 0.0 2.0 0.0 0.0
 0.0 0.0 3.0 0.0
 0.0 0.0 0.0 3.0

```

- In general, it is expected for `gradient` to preserve the structure and for `jacobian` to forget it. A gradient lives in the same space as the original input, whereas a Jacobian lives in a kind of product space. Even in the simple case of a matrix-to-matrix function, most autodiff backends (except Enzyme I think) will flatten it and represent its Jacobian as a 2d matrix instead of a 4d tensor.

- Theoretically we could want a `ComponentMatrix` as the result but given how autodiff works, the Jacobian is usually either (1) a horizontal stacking of output-like vectors or (2) a vertical stacking of input-like vectors. It is difficult to take both input and output structure into account at the same time. Also, `ComponentMatrix` was kind of [buggy](https://github.com/SciML/ComponentArrays.jl/issues/258) last time I checked.

- If you want to use sparse autodiff with DI, then you shouldn’t want a `ComponentMatrix` anyway, but a `SparseMatrixCSC` or some kind of block matrix, e.g. with BlockArrays.jl?

- What do you mean by “preserve the names for computing chain rules”? What is the point of autodiff if you compute chain rules manually?

---

<div class="post-metadata">

**Author:** ![charlesknipp](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/charlesknipp/32/214437_2.png) [@charlesknipp](https://discourse.julialang.org/u/charlesknipp)\
**Post date:** [October 23, 2025, 6:23pm UTC](https://discourse.julialang.org/t/preserving-componentarray-structure-with-differentiationinterface/133390/3 "2025-10-23T18:23:52Z")

</div>

> [@gdalle](#):
>
> A gradient lives in the same space as the original input, whereas a Jacobian lives in a kind of product space.

Ahh yes, that makes quite a bit of sense in this case. I guess I may need to reshape the output since it at least follows a consistent structure.

> [@gdalle](#):
>
> What do you mean by “preserve the names for computing chain rules”? What is the point of autodiff if you compute chain rules manually?

Another paradigm of writing macro models is based on a DAG, where the various blocks (nodes) represent a group of model equations. You can reduce the dimension of the problem when computing block Jacobians and chaining ruling them together (as opposed to putting everything in one block). For large models with heterogeneous agents, solving models is quite a bit faster under this approach. [This](http://mattrognlie.com/sequence_space_jacobian.pdf) seminal macro paper illustrates this concept further.

> [@gdalle](#):
>
> If you want to use sparse autodiff with DI, then you shouldn’t want a `ComponentMatrix` anyway, but a `SparseMatrixCSC` or some kind of block matrix, e.g. with [BlockArrays.jl](https://juliaregistries.github.io/General/packages/redirect_to_repo/BlockArrays)?

This is definitely the proper approach, and one I’m likely to take in the future.

Thanks for your work on DI btw, it has been an amazing addition to my everyday carry.