diff --git a/src/trace_functions.jl b/src/trace_functions.jl index d678119..56e2f0e 100644 --- a/src/trace_functions.jl +++ b/src/trace_functions.jl @@ -49,9 +49,18 @@ function trace_function(::Type{T}, f, x) where {T<:Union{AbstractTracer,Dual}} return xt, yt end -function trace_function(::Type{T}, f!, y, x) where {T<:Union{AbstractTracer,Dual}} +function trace_function(::Type{T}, f!, y, x) where {T<:AbstractTracer} xt = trace_input(T, x) yt = similar(y, T) + fill!(yt, myempty(T)) + f!(yt, xt) + return xt, yt +end + +function trace_function(::Type{D}, f!, y, x) where {P,T<:AbstractTracer,D<:Dual{P,T}} + t = myempty(T) + xt = trace_input(D, x) + yt = Dual.(y, t) f!(yt, xt) return xt, yt end diff --git a/test/componentarrays.jl b/test/componentarrays.jl index a341bb9..b8e0e3b 100644 --- a/test/componentarrays.jl +++ b/test/componentarrays.jl @@ -21,10 +21,11 @@ y_comp = ComponentVector(; a=rand(2), b=rand(3)) x = Vector(x_comp) y = Vector(y_comp) -detector = TracerSparsityDetector() - -@test jacobian_sparsity(f_comp, x_comp, detector) == jacobian_sparsity(f, x, detector) -@test jacobian_sparsity(f_comp!, similar(y_comp), x_comp, detector) == - jacobian_sparsity(f!, similar(y), x, detector) -@test hessian_sparsity(sum ∘ f_comp, x_comp, detector) == - hessian_sparsity(sum ∘ f, x, detector) +@testset "$detector" for detector in + (TracerSparsityDetector(), TracerLocalSparsityDetector()) + @test jacobian_sparsity(f_comp, x_comp, detector) == jacobian_sparsity(f, x, detector) + @test jacobian_sparsity(f_comp!, similar(y_comp), x_comp, detector) == + jacobian_sparsity(f!, similar(y), x, detector) + @test hessian_sparsity(sum ∘ f_comp, x_comp, detector) == + hessian_sparsity(sum ∘ f, x, detector) +end diff --git a/test/test_gradient.jl b/test/test_gradient.jl index 87dc127..02dbb8a 100644 --- a/test/test_gradient.jl +++ b/test/test_gradient.jl @@ -14,12 +14,16 @@ REAL_TYPES = (Float64, Int, Bool, UInt8, Float16, Rational{Int}) # NOTE: J gets overwritten inside the testsets. detector = TracerSparsityDetector() J(f, x) = jacobian_sparsity(f, x, detector) +J(f!, y, x) = jacobian_sparsity(f!, y, x, detector) +P = first(GRADIENT_PATTERNS) +T = GradientTracer{P} @testset "Jacobian Global" begin @testset "$P" for P in GRADIENT_PATTERNS T = GradientTracer{P} detector = TracerSparsityDetector(; gradient_tracer_type=T) J(f, x) = jacobian_sparsity(f, x, detector) + J(f!, y, x) = jacobian_sparsity(f!, y, x, detector) @testset "Trivial examples" begin f(x) = [x[1]^2, 2 * x[1] * x[2]^2, sin(x[3])] @@ -53,6 +57,18 @@ J(f, x) = jacobian_sparsity(f, x, detector) @test J(x -> ℯ^zero(x), 1) ≈ [0;;] end + @testset "In-place functions" begin + x = rand(5) + y = similar(x) + + function f!(y, x) + for i in 1:(length(x) - 1) + y[i] = x[i + 1] - x[i] + end + end + @test_nowarn J(f!, y, x) + end + # Conversions @testset "Conversion" begin @testset "to $T" for T in REAL_TYPES @@ -162,6 +178,7 @@ end T = GradientTracer{P} detector = TracerLocalSparsityDetector(; gradient_tracer_type=T) J(f, x) = jacobian_sparsity(f, x, detector) + J(f!, y, x) = jacobian_sparsity(f!, y, x, detector) @testset "Trivial examples" begin @@ -242,6 +259,18 @@ end @test J(x -> 0, 1) ≈ [0;;] end + @testset "In-place functions" begin + x = rand(5) + y = similar(x) + + function f!(y, x) + for i in 1:(length(x) - 1) + y[i] = x[i + 1] - x[i] + end + end + @test_nowarn J(f!, y, x) + end + # Conversions @testset "Conversion" begin @testset "Conversion to $T" for T in REAL_TYPES