Reverse Jacobian on StaticArrays

[gdalle]

Apparently reverse-mode jacobian tries to mutate something, and StaticArrays aren't happy:

julia> import Enzyme

julia> using StaticArrays

julia> x = @SVector(ones(2))
2-element SVector{2, Float64} with indices SOneTo(2):
 1.0
 1.0

julia> Enzyme.jacobian(Enzyme.Reverse, identity, x)
ERROR: setindex!(::SVector{2, Float64}, value, ::Int) is not defined.
 Hint: Use `MArray` or `SizedArray` to create a mutable static array
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] setindex!(a::SVector{2, Float64}, value::Float64, i::Int64)
    @ StaticArrays ~/.julia/packages/StaticArrays/lZeGV/src/indexing.jl:3
  [3] zerosetfn!(x::SVector{2, Float64}, i::Int64, val::Float64)
    @ Enzyme ~/.julia/packages/Enzyme/R6sE8/src/sugar.jl:11
  [4] #130
    @ ~/.julia/packages/Enzyme/R6sE8/src/sugar.jl:935 [inlined]
  [5] ntuple
    @ ./ntuple.jl:49 [inlined]
  [6] jacobian(mode::EnzymeCore.ReverseMode{…}, f::typeof(identity), x::SVector{…}; n_outs::Val{…}, chunk::Nothing)
    @ Enzyme ~/.julia/packages/Enzyme/R6sE8/src/sugar.jl:929
  [7] jacobian
    @ ~/.julia/packages/Enzyme/R6sE8/src/sugar.jl:846 [inlined]
  [8] #jacobian#129
    @ ~/.julia/packages/Enzyme/R6sE8/src/sugar.jl:861 [inlined]
  [9] jacobian(mode::EnzymeCore.ReverseMode{…}, f::typeof(identity), x::SVector{…})
    @ Enzyme ~/.julia/packages/Enzyme/R6sE8/src/sugar.jl:846
 [10] top-level scope
    @ ~/Documents/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterface/test/Back/Enzyme/playground.jl:8
Some type information was truncated. Use `show(err)` to see complete types.

[alpttex19]

have you solved this problem?

[gdalle]

Due to this mutation issue, the reverse-mode Jacobian in DifferentiationInterface avoids Enzyme's native API for now. As a result, it works correctly for this case and makes no allocations once preparation is done:

julia> using DifferentiationInterface, BenchmarkTools; import Enzyme

julia> backend = AutoEnzyme(mode=Enzyme.Reverse);

julia> prep = prepare_jacobian(identity, backend, x);

julia> jacobian(identity, prep, backend, x)
2×2 SMatrix{2, 2, Float64, 4} with indices SOneTo(2)×SOneTo(2):
 1.0  0.0
 0.0  1.0

julia> @btime jacobian($identity, $prep, $backend, $x);
  3.416 ns (0 allocations: 0 bytes)

But I don't think the fix in Enzyme would be too hard, I'll defer to the maintainers