Type stability of gradient and jacobian

[KnutAM]

I couldn't find any issues on this, but noticed that the return type of gradient and jacobian is not type stable without passing the config (or using a mutating function in the jacobian case).

Took some time to track this down today as it was nested deep, of course just pre-allocating the output and using the jacobian! function makes it fine again. Has there been some discussions on how to make this stable?

julia> @code_warntype ForwardDiff.gradient(sum, [1.])
MethodInstance for ForwardDiff.gradient(::typeof(sum), ::Vector{Float64})
  from gradient(f, x::AbstractArray) @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:16
Arguments
  #self#::Core.Const(ForwardDiff.gradient)
  f::Core.Const(sum)
  x::Vector{Float64}
Body::Any
1 ─ %1 = ForwardDiff.GradientConfig(f, x)::ForwardDiff.GradientConfig{ForwardDiff.Tag{typeof(sum), Float64}, Float64, _A, Array{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sum), Float64}, Float64, _A1}, 1}} where {_A, _A1}
│   %2 = (#self#)(f, x, %1)::Any
└──      return %2

julia> @code_warntype ForwardDiff.jacobian(identity, [1.0])
MethodInstance for ForwardDiff.jacobian(::typeof(identity), ::Vector{Float64})
  from jacobian(f, x::AbstractArray) @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:18
Arguments
  #self#::Core.Const(ForwardDiff.jacobian)
  f::Core.Const(identity)
  x::Vector{Float64}
Body::Union{Matrix{Any}, Matrix{Float64}, Matrix{E} where E<:(ForwardDiff.Dual{ForwardDiff.Tag{typeof(identity), Float64}, Float64})}
1 ─ %1 = ForwardDiff.JacobianConfig(f, x)::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(identity), Float64}, Float64, _A, Array{ForwardDiff.Dual{ForwardDiff.Tag{typeof(identity), Float64}, Float64, _A1}, 1}} where {_A, _A1}
│   %2 = (#self#)(f, x, %1)::Union{Matrix{Any}, Matrix{Float64}, Matrix{E} where E<:(ForwardDiff.Dual{ForwardDiff.Tag{typeof(identity), Float64}, Float64})}
└──      return %2

Either passing cfg = JacobianConfig(...) or using a mutating residual in the jacobian case, solves the problem, i.e.

julia> @code_warntype ForwardDiff.jacobian((y,x) -> (y .= x), rand(1), [1.0])
MethodInstance for ForwardDiff.jacobian(::var"#3#4", ::Vector{Float64}, ::Vector{Float64})
  from jacobian(f!, y::AbstractArray, x::AbstractArray) @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:35
Arguments
  #self#::Core.Const(ForwardDiff.jacobian)
  f!::Core.Const(var"#3#4"())
  y::Vector{Float64}
  x::Vector{Float64}
Body::Matrix{Float64}
1 ─ %1 = ForwardDiff.JacobianConfig(f!, y, x)::ForwardDiff.JacobianConfig{ForwardDiff.Tag{var"#3#4", Float64}, Float64, _A, <:Tuple{Array{ForwardDiff.Dual{ForwardDiff.Tag{var"#3#4", Float64}, Float64, _A}, 1} where _A, Array{ForwardDiff.Dual{ForwardDiff.Tag{var"#3#4", Float64}, Float64, _A}, 1} where _A}} where _A
│   %2 = (#self#)(f!, y, x, %1)::Matrix{Float64}
└──      return %2

The only way I found to solve this, is limiting (at least in these cases when the cfg is automatically created leading to a type instability) the chunk size, e.g.

const ChunkUnions3 = Union{(ForwardDiff.Chunk{i} for i in (2, 4, 12))...}
function ForwardDiff.Chunk(input_length::Integer, threshold::Integer = ForwardDiff.DEFAULT_CHUNK_THRESHOLD)
    N = ForwardDiff.pickchunksize(input_length, threshold)
    chunk = Base.@nif 12 d->(N == d) d->(Chunk{d}()) d->(Chunk{N}())
    return chunk::ChunkUnions3
end

julia> @code_warntype ForwardDiff.jacobian(identity, [1.0])
MethodInstance for ForwardDiff.jacobian(::typeof(identity), ::Vector{Float64})
  from jacobian(f::F, x::AbstractArray) where F @ ForwardDiff ~/.julia/packages/ForwardDiff/X74OO/src/jacobian.jl:18
Static Parameters
  F = typeof(identity)
Arguments
  #self#::Core.Const(ForwardDiff.jacobian)
  f::Core.Const(identity)
  x::Vector{Float64}
Body::Matrix{Float64}
1 ─ %1 = ForwardDiff.JacobianConfig(f, x)::Union{ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(identity), Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(identity), Float64}, Float64, 2}}}, ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(identity), Float64}, Float64, 4, Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(identity), Float64}, Float64, 4}}}, ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(identity), Float64}, Float64, 12, Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(identity), Float64}, Float64, 12}}}}
│   %2 = (#self#)(f, x, %1)::Matrix{Float64}
└──      return %2

[devmotion]

This is known and documented:

If your input dimension is constant across calls, you should explicitly select a chunk size rather than relying on ForwardDiff's heuristic. There are two reasons for this. The first is that ForwardDiff's heuristic depends only on the input dimension, whereas in reality the optimal chunk size will also depend on the target function. The second is that ForwardDiff's heuristic is inherently type-unstable, which can cause the entire call to be type-unstable.

If your input dimension is a runtime variable, you can rely on ForwardDiff's heuristic without sacrificing type stability by manually asserting the output type, or - even better - by using the in-place API functions:

Another way to avoid this instability if you compute gradients, Jacobians etc. repeatedly would be to use the DifferentiationInterface API with preparation (which implicitly creates a config object for you to be reused in subsequent gradient etc calculations).

[KristofferC]

Even if the chunksize is unknown it should still be possible to infer the output type. See #558 for a similar fixed issue.

[KnutAM]

This is known and documented:

Thanks - just searched for issues related to this!

I missed this as usually I use the inplace/preallicate the chunks, but realized that for new users this can really slow down and be hard to track if you don't know what to look for.

But as I understand- you wouldn't accept a slower execution even if it would solve this issue?

Even if the chunksize is unknown it should still be possible to infer the output type. See #558 for a similar fixed issue.

That's what I thought would happen, even if the diff call itself is unstable. But I guess the problem for the out of place is knowing the return type of f(x) for x nondual without calling it? In my example above, adding more sizes to the Union made it fail...