diff --git a/Project.toml b/Project.toml index 84c81511..77f40590 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ForwardDiff" uuid = "f6369f11-7733-5829-9624-2563aa707210" -version = "1.4.0" +version = "1.4.1" [deps] CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950" diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index ec8cbc83..bf0ef99a 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -21,8 +21,6 @@ using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult end end -@inline static_dual_eval(::Type{T}, f::F, x::StaticArray) where {T,F} = f(dualize(T, x)) - # To fix method ambiguity issues: function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) where {Tg,T<:Real,N} return ForwardDiff._eigvals(A) @@ -54,12 +52,12 @@ end @inline function ForwardDiff.vector_mode_gradient(f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - return extract_gradient(T, static_dual_eval(T, f, x), x) + return extract_gradient(T, f(dualize(T, x)), x) end @inline function ForwardDiff.vector_mode_gradient!(result, f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - return extract_gradient!(T, result, static_dual_eval(T, f, x)) + return extract_gradient!(T, result, f(dualize(T, x))) end # Jacobian @@ -84,7 +82,7 @@ end @inline function ForwardDiff.vector_mode_jacobian(f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - return extract_jacobian(T, static_dual_eval(T, f, x), x) + return extract_jacobian(T, f(dualize(T, x)), x) end function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T @@ -94,7 +92,7 @@ end @inline function ForwardDiff.vector_mode_jacobian!(result, f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - ydual = static_dual_eval(T, f, x) + ydual = f(dualize(T, x)) result = extract_jacobian!(T, result, ydual, length(x)) result = extract_value!(T, result, ydual) return result @@ -102,7 +100,7 @@ end @inline function ForwardDiff.vector_mode_jacobian!(result::ImmutableDiffResult, f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - ydual = static_dual_eval(T, f, x) + ydual = f(dualize(T, x)) result = DiffResults.jacobian!(result, extract_jacobian(T, ydual, x)) result = DiffResults.value!(Base.Fix1(value, T), result, ydual) return result diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 36df8eba..af8d6e77 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -1,6 +1,7 @@ module AllocationsTest using ForwardDiff +using StaticArrays include(joinpath(dirname(@__FILE__), "utils.jl")) @@ -44,4 +45,25 @@ end @test iszero(allocs_jacobian!()) end +@testset "allocation-free nested StaticArray jacobian" begin + # test that nested jacobians of StaticArrays do not allocate. + # This is a regression test for issue #798, where the inner jacobian was allocating + toy_f(x) = SVector(x[1]^2 * x[2], sin(x[1]) + x[2]^3) + + function toy_J_flat(x::SVector{2,T}) where {T} + y = ForwardDiff.jacobian(toy_f, x) + return SVector{2,T}(y[1], y[2]) + end + + function toy_nested_jacobian(x0::SVector{2,T}) where {T} + return ForwardDiff.jacobian(toy_J_flat, x0) + end + + function allocs_jacobian() + x0 = SVector(1.0, 2.0) + return @allocated toy_nested_jacobian(x0) + end + @test iszero(allocs_jacobian()) +end + end