value_and_jacobian inference problem with AutoForwardDiff

[tpapp]

MWE:

julia> using StaticArrays, DifferentiationInterface, ForwardDiff, Test

julia> const A = one(SMatrix{3, 3})
3×3 SMatrix{3, 3, Float64, 9} with indices SOneTo(3)×SOneTo(3):
 1.0  0.0  0.0
 0.0  1.0  0.0
 0.0  0.0  1.0

julia> f(x) = A * x
f (generic function with 1 method)

julia> z1 = zeros(3)
3-element Vector{Float64}:
 0.0
 0.0
 0.0

julia> z2 = zero(SVector{3})
3-element SVector{3, Float64} with indices SOneTo(3):
 0.0
 0.0
 0.0

julia> @inferred value_and_jacobian(f, AutoForwardDiff(), z1)
ERROR: return type Tuple{SVector{3, Float64}, Matrix{Float64}} does not match inferred return type Tuple{SVector{3, Float64}, Any}
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] top-level scope
   @ REPL[18]:1

julia> @inferred value_and_jacobian(f, AutoForwardDiff(), z2) # works fine, included for comparison
([0.0, 0.0, 0.0], [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0])

julia> VERSION
v"1.11.1"

(tmp) pkg> st
Status `/tmp/Project.toml`
  [a0c0ee7d] DifferentiationInterface v0.6.23 `https://github.com/JuliaDiff/DifferentiationInterface.jl.git:DifferentiationInterface#main`
  [f6369f11] ForwardDiff v0.10.38
  [90137ffa] StaticArrays v1.9.8

I included an SVector for comparison. Looking into this, I think the type unstable part is the tagging.

[gdalle]

Unfortunately, it is a well-known limitation of ForwardDiff itself: the dynamic choice of chunk size leads to type instability. You don't actually need DI to see this:

julia> using Test; import ForwardDiff

julia> @inferred ForwardDiff.jacobian(identity, [1.0])
ERROR: return type Matrix{Float64} does not match inferred return type Union{Matrix{Any}, Matrix{Float64}, Matrix{E} where E<:(ForwardDiff.Dual{ForwardDiff.Tag{typeof(identity), Float64}, Float64})}

However, DI offers you two workarounds. You can either pick a backend with a fixed chunk size (ensuring that it is always smaller than the input dimension):

julia> import DifferentiationInterface as DI

julia> @inferred DI.jacobian(identity, DI.AutoForwardDiff(; chunksize=1), [1.0])
1×1 Matrix{Float64}:
 1.0

Or you can keep the adaptive backend but use the preparation mechanism. In the case of ForwardDiff, DI's preparation includes chunk size selection. It is type-unstable but you only need to do it once:

julia> prep = DI.prepare_jacobian(identity, DI.AutoForwardDiff(), [0.0]);

julia> @inferred DI.jacobian(identity, prep, DI.AutoForwardDiff(), [0.0])
1×1 Matrix{Float64}:
 1.0

Does that solve your problem?

[tpapp]

Using the preparation mechanism just shifts the type instability elsewhere.

ensuring that it is always smaller than the input dimension

AFAIK it is not necessary, ForwardDiff works fine when the chunk size is above the input length (may not be optimal of course). No, that was a misunderstanding on my part, it would require a PR to ForwardDiff to fall back to that.

FWIW, I would just get rid of the heuristic and pick a not-too-large fixed chunk size as the default, eg 4 or 8, and suggest that the user modifies it if desired. That leaves 1 as the default chunk size.

[gdalle]

Using the preparation mechanism just shifts the type instability elsewhere.

That is true, but the important thing is that you can shift it outside of the hot loop. If you compute 1000 Jacobians of the same function on inputs of the same type and size, you only need to prepare once.

AFAIK it is not necessary, ForwardDiff works fine when the chunk size is above the input length (may not be optimal of course).

That is also true but the whole point of fixing the chunk size is to guarantee type stability of Jacobian config creation. If you keep the step where you threshold the chunk size using the input size, you still have a type-unstable mechanism.
See also the discussion in #539 where I introduced this change.

FWIW, I would just get rid of the heuristic and pick a not-too-large fixed chunk size as the default, eg 4 or 8, and suggest that the user modifies it if desired.

That's a change which would need to be done in ForwardDiff itself. There have been discussions around it, but until the type-unstable chunk size selection is removed, DI can't really be smarter than ForwardDiff in this regard.

[tpapp]

Thanks for the explanation, closing this.