Relax type annotations in Jacobian output parsing (#217)

* Move output parsing into separate file

* Rename functions

* Loosen type restrictions

* Add tests

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseConnectivityTracer"
 uuid = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 authors = ["Adrian Hill <[email protected]>"]
-version = "0.6.8"
+version = "0.6.9-DEV"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/SparseConnectivityTracer.jl

@@ -31,6 +31,7 @@ include("overloads/arrays.jl")
 include("overloads/ambiguities.jl")
 
 include("trace_functions.jl")
+include("parse_outputs_to_matrix.jl")
 include("adtypes_interface.jl")
 
 export TracerSparsityDetector

src/parse_outputs_to_matrix.jl

@@ -0,0 +1,76 @@
+#= Parse tracers from function evaluation in `src/trace_functions.jl` into an output matrix. =#
+
+_tracer_or_number(x::Real) = x
+_tracer_or_number(d::Dual) = tracer(d)
+
+#==========#
+# Jacobian #
+#==========#
+
+function jacobian_tracers_to_matrix(
+    xt::AbstractArray{T}, yt::AbstractArray
+) where {T<:GradientTracer}
+    n, m = length(xt), length(yt)
+    I = Int[] # row indices
+    J = Int[] # column indices
+    V = Bool[]   # values
+    for (i, y) in enumerate(yt)
+        if y isa T && !isemptytracer(y)
+            for j in gradient(y)
+                push!(I, i)
+                push!(J, j)
+                push!(V, true)
+            end
+        end
+    end
+    return sparse(I, J, V, m, n)
+end
+
+function jacobian_tracers_to_matrix(
+    xt::AbstractArray{D}, yt::AbstractArray
+) where {P,T<:GradientTracer,D<:Dual{P,T}}
+    return jacobian_tracers_to_matrix(tracer.(xt), _tracer_or_number.(yt))
+end
+
+#=========#
+# Hessian #
+#=========#
+
+function hessian_tracers_to_matrix(xt::AbstractArray{T}, yt::T) where {T<:HessianTracer}
+    n = length(xt)
+    I = Int[] # row indices
+    J = Int[] # column indices
+    V = Bool[]   # values
+
+    if !isemptytracer(yt)
+        for (i, j) in tuple_set(hessian(yt))
+            push!(I, i)
+            push!(J, j)
+            push!(V, true)
+            # TODO: return `Symmetric` instead on next breaking release
+            push!(I, j)
+            push!(J, i)
+            push!(V, true)
+        end
+    end
+    h = sparse(I, J, V, n, n)
+    return h
+end
+
+function hessian_tracers_to_matrix(
+    xt::AbstractArray{D1}, yt::D2
+) where {P1,P2,T<:HessianTracer,D1<:Dual{P1,T},D2<:Dual{P2,T}}
+    return hessian_tracers_to_matrix(tracer.(xt), tracer(yt))
+end
+
+function hessian_tracers_to_matrix(
+    xt::AbstractArray{T}, yt::Number
+) where {T<:HessianTracer}
+    return hessian_tracers_to_matrix(xt, myempty(T))
+end
+
+function hessian_tracers_to_matrix(
+    xt::AbstractArray{D1}, yt::Number
+) where {P1,T<:HessianTracer,D1<:Dual{P1,T}}
+    return hessian_tracers_to_matrix(tracer.(xt), myempty(T))
+end

src/trace_functions.jl

@@ -1,7 +1,8 @@
 #= This file handles the actual tracing of functions:
 1) creating tracers from inputs
 2) evaluating the function with the created tracers
-3) parsing the resulting tracers into an output matrix
+
+The resulting output is parsed in `src/parse_outputs_to_matrix.jl`.
 =#
 
 #==================#
@@ -58,28 +59,24 @@ end
 to_array(x::Real) = [x]
 to_array(x::AbstractArray) = x
 
-# Utilities
-_tracer_or_number(x::Real) = x
-_tracer_or_number(d::Dual) = tracer(d)
-
-#================#
-# GradientTracer #
-#================#
+#==========#
+# Jacobian #
+#==========#
 
 # Compute the sparsity pattern of the Jacobian of `y = f(x)`.
 function _jacobian_sparsity(
     f, x, ::Type{T}=DEFAULT_GRADIENT_TRACER
 ) where {T<:GradientTracer}
     xt, yt = trace_function(T, f, x)
-    return jacobian_pattern_to_mat(to_array(xt), to_array(yt))
+    return jacobian_tracers_to_matrix(to_array(xt), to_array(yt))
 end
 
 # Compute the sparsity pattern of the Jacobian of `f!(y, x)`.
 function _jacobian_sparsity(
     f!, y, x, ::Type{T}=DEFAULT_GRADIENT_TRACER
 ) where {T<:GradientTracer}
     xt, yt = trace_function(T, f!, y, x)
-    return jacobian_pattern_to_mat(to_array(xt), to_array(yt))
+    return jacobian_tracers_to_matrix(to_array(xt), to_array(yt))
 end
 
 # Compute the local sparsity pattern of the Jacobian of `y = f(x)` at `x`.
@@ -88,7 +85,7 @@ function _local_jacobian_sparsity(
 ) where {T<:GradientTracer}
     D = Dual{eltype(x),T}
     xt, yt = trace_function(D, f, x)
-    return jacobian_pattern_to_mat(to_array(xt), to_array(yt))
+    return jacobian_tracers_to_matrix(to_array(xt), to_array(yt))
 end
 
 # Compute the local sparsity pattern of the Jacobian of `f!(y, x)` at `x`.
@@ -97,42 +94,17 @@ function _local_jacobian_sparsity(
 ) where {T<:GradientTracer}
     D = Dual{eltype(x),T}
     xt, yt = trace_function(D, f!, y, x)
-    return jacobian_pattern_to_mat(to_array(xt), to_array(yt))
-end
-
-function jacobian_pattern_to_mat(
-    xt::AbstractArray{T}, yt::AbstractArray{<:Real}
-) where {T<:GradientTracer}
-    n, m = length(xt), length(yt)
-    I = Int[] # row indices
-    J = Int[] # column indices
-    V = Bool[]   # values
-    for (i, y) in enumerate(yt)
-        if y isa T && !isemptytracer(y)
-            for j in gradient(y)
-                push!(I, i)
-                push!(J, j)
-                push!(V, true)
-            end
-        end
-    end
-    return sparse(I, J, V, m, n)
-end
-
-function jacobian_pattern_to_mat(
-    xt::AbstractArray{D}, yt::AbstractArray{<:Real}
-) where {P,T<:GradientTracer,D<:Dual{P,T}}
-    return jacobian_pattern_to_mat(tracer.(xt), _tracer_or_number.(yt))
+    return jacobian_tracers_to_matrix(to_array(xt), to_array(yt))
 end
 
-#===============#
-# HessianTracer #
-#===============#
+#=========#
+# Hessian #
+#=========#
 
 # Compute the sparsity pattern of the Hessian of a scalar function `y = f(x)`.
 function _hessian_sparsity(f, x, ::Type{T}=DEFAULT_HESSIAN_TRACER) where {T<:HessianTracer}
     xt, yt = trace_function(T, f, x)
-    return hessian_pattern_to_mat(to_array(xt), yt)
+    return hessian_tracers_to_matrix(to_array(xt), yt)
 end
 
 # Compute the local sparsity pattern of the Hessian of a scalar function `y = f(x)` at `x`.
@@ -141,42 +113,5 @@ function _local_hessian_sparsity(
 ) where {T<:HessianTracer}
     D = Dual{eltype(x),T}
     xt, yt = trace_function(D, f, x)
-    return hessian_pattern_to_mat(to_array(xt), yt)
-end
-
-function hessian_pattern_to_mat(xt::AbstractArray{T}, yt::T) where {T<:HessianTracer}
-    n = length(xt)
-    I = Int[] # row indices
-    J = Int[] # column indices
-    V = Bool[]   # values
-
-    if !isemptytracer(yt)
-        for (i, j) in tuple_set(hessian(yt))
-            push!(I, i)
-            push!(J, j)
-            push!(V, true)
-            # TODO: return `Symmetric` instead on next breaking release
-            push!(I, j)
-            push!(J, i)
-            push!(V, true)
-        end
-    end
-    h = sparse(I, J, V, n, n)
-    return h
-end
-
-function hessian_pattern_to_mat(
-    xt::AbstractArray{D1}, yt::D2
-) where {P1,P2,T<:HessianTracer,D1<:Dual{P1,T},D2<:Dual{P2,T}}
-    return hessian_pattern_to_mat(tracer.(xt), tracer(yt))
-end
-
-function hessian_pattern_to_mat(xt::AbstractArray{T}, yt::Number) where {T<:HessianTracer}
-    return hessian_pattern_to_mat(xt, myempty(T))
-end
-
-function hessian_pattern_to_mat(
-    xt::AbstractArray{D1}, yt::Number
-) where {P1,T<:HessianTracer,D1<:Dual{P1,T}}
-    return hessian_pattern_to_mat(tracer.(xt), myempty(T))
+    return hessian_tracers_to_matrix(to_array(xt), yt)
 end

test/test_gradient.jl

@@ -123,6 +123,36 @@ J(f, x) = jacobian_sparsity(f, x, detector)
                 x -> x[1] > x[2] ? x[3] : x[4], [1.0, 2.0, 3.0, 4.0]
             ) == [0 0 1 1;]
         end
+
+        @testset "Output of type Vector{Any}" begin
+            function f(x::AbstractVector)
+                n = length(x)
+                ret = [] # return type will be Vector{Any}
+                for i in 1:(n - 1)
+                    append!(
+                        ret,
+                        abs2(x[i + 1]) - abs2(x[i]) + abs2(x[n - i]) - abs2(x[n - i + 1]),
+                    )
+                end
+                return ret
+            end
+            x = [
+                0.263914
+                0.605532
+                1.281598
+                1.413663
+                0.178133
+                -1.705427
+            ]
+            @test J(f, x) == [
+                1  1  0  0  1  1
+                0  1  1  1  1  0
+                0  0  1  1  0  0
+                0  1  1  1  1  0
+                1  1  0  0  1  1
+            ]
+        end
+
         yield()
     end
 end
@@ -248,6 +278,35 @@ end
             @test J(x -> log(det(x)), [1.0 -1.0; 2.0 2.0]) == [1 1 1 1]
             @test J(x -> dot(x[1:2], x[4:5]), [0, 1, 0, 1, 0]) == [1 0 0 0 1]
         end
+        @testset "Output of type Vector{Any}" begin
+            function f(x::AbstractVector)
+                n = length(x)
+                ret = [] # return type will be Vector{Any}
+                for i in 1:(n - 1)
+                    append!(
+                        ret,
+                        abs2(x[i + 1]) - abs2(x[i]) + abs2(x[n - i]) - abs2(x[n - i + 1]),
+                    )
+                end
+                return ret
+            end
+            x = [
+                0.263914
+                0.605532
+                1.281598
+                1.413663
+                0.178133
+                -1.705427
+            ]
+            @test J(f, x) == [
+                1  1  0  0  1  1
+                0  1  1  1  1  0
+                0  0  1  1  0  0
+                0  1  1  1  1  0
+                1  1  0  0  1  1
+            ]
+        end
+
         yield()
     end
 end