Support ComponentArrays

Project.toml

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

src/interface.jl

@@ -17,7 +17,8 @@ Supports [`GradientTracer`](@ref), [`HessianTracer`](@ref) and [`Dual`](@ref).
 trace_input(::Type{T}, xs) where {T<:Union{AbstractTracer,Dual}} = trace_input(T, xs, 1)
 
 function trace_input(::Type{T}, xs::AbstractArray, i) where {T<:Union{AbstractTracer,Dual}}
-    is = reshape(1:length(xs), size(xs)) .+ (i - 1)
+    is = similar(xs, Int)  # same array type as xs
+    is .= reshape(1:length(xs), size(xs)) .+ (i - 1)
     return create_tracers(T, xs, is)
 end
 function trace_input(::Type{T}, x::Real, i::Integer) where {T<:Union{AbstractTracer,Dual}}

test/Project.toml

@@ -4,6 +4,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"

test/componentarrays.jl

@@ -0,0 +1,30 @@
+using ComponentArrays
+using SparseConnectivityTracer
+using Test
+
+f(x::AbstractVector) = abs2.(x)
+f_comp(x::ComponentVector) = ComponentVector(; a=abs2.(x.a), b=abs2.(x.b))
+
+function f!(y::AbstractVector, x::AbstractVector)
+    y .= abs2.(x)
+    return y
+end
+
+function f_comp!(y::ComponentVector, x::ComponentVector)
+    y.a .= abs2.(x.a)
+    y.b .= abs2.(x.b)
+    return y
+end
+
+x_comp = ComponentVector(; a=rand(2), b=rand(3))
+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)

test/runtests.jl

@@ -90,6 +90,9 @@ GROUP = get(ENV, "JULIA_SCT_TEST_GROUP", "Core")
             @testset "Array overloads" begin
                 include("test_arrays.jl")
             end
+            @testset "ComponentArrays" begin
+                include("componentarrays.jl")
+            end
         end
     end