Fix jacobian_sparsity for inplace function (#223)

* Fix jacobian_sparsity for inplace function * Test both local and global inplace functions * Fix initialization of `Dual` tracers * Test ComponentArrays globally and locally --------- Co-authored-by: adrhill <adrian.hill@mailbox.org>
adrhill · Jan 24, 2025 · 46d9aac · 46d9aac
1 parent 54f7c97
commit 46d9aac
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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