feat: support wrapper types as well

Author: avik-pal

Project.toml

@@ -15,11 +15,9 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [weakdeps]
-GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [extensions]
-ForwardDiffGPUArraysCoreExt = "GPUArraysCore"
 ForwardDiffStaticArraysExt = "StaticArrays"
 
 [compat]
@@ -28,7 +26,6 @@ CommonSubexpressions = "0.3"
 DiffResults = "1.1"
 DiffRules = "1.4"
 DiffTests = "0.1"
-GPUArraysCore = "0.2"
 IrrationalConstants = "0.1, 0.2"
 LogExpFunctions = "0.3"
 NaNMath = "1"

ext/ForwardDiffGPUArraysCoreExt.jl

@@ -21,6 +21,7 @@ include("derivative.jl")
 include("gradient.jl")
 include("jacobian.jl")
 include("hessian.jl")
+include("utils.jl")
 
 export DiffResults
 

src/ForwardDiff.jl

@@ -72,36 +72,46 @@ end
 
 function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
                seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
-    if isbitstype(V)
-        for idx in structural_eachindex(duals, x)
-            duals[idx] = Dual{T,V,N}(x[idx], seed)
-        end
-    else
-        for idx in structural_eachindex(duals, x)
-            if isassigned(x, idx)
+    if supports_fast_scalar_indexing(duals)
+        if isbitstype(V)
+            for idx in structural_eachindex(duals, x)
                 duals[idx] = Dual{T,V,N}(x[idx], seed)
-            else
-                Base._unsetindex!(duals, idx)
+            end
+        else
+            for idx in structural_eachindex(duals, x)
+                if isassigned(x, idx)
+                    duals[idx] = Dual{T,V,N}(x[idx], seed)
+                else
+                    Base._unsetindex!(duals, idx)
+                end
             end
         end
+    else
+        idxs = collect(structural_eachindex(duals, x))
+        duals[idxs] .= Dual{T,V,N}.(view(x, idxs), Ref(seed))
     end
     return duals
 end
 
 function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
                seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
-    if isbitstype(V)
-        for (i, idx) in zip(1:N, structural_eachindex(duals, x))
-            duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
-        end
-    else
-        for (i, idx) in zip(1:N, structural_eachindex(duals, x))
-            if isassigned(x, idx)
+    if supports_fast_scalar_indexing(duals)
+        if isbitstype(V)
+            for (i, idx) in zip(1:N, structural_eachindex(duals, x))
                 duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
-            else
-                Base._unsetindex!(duals, idx)
+            end
+        else
+            for (i, idx) in zip(1:N, structural_eachindex(duals, x))
+                if isassigned(x, idx)
+                    duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
+                else
+                    Base._unsetindex!(duals, idx)
+                end
             end
         end
+    else
+        idxs = collect(Iterators.take(structural_eachindex(duals, x), N))
+        duals[idxs] .= Dual{T,V,N}.(view(x, idxs), getindex.(Ref(seeds), 1:length(idxs)))
     end
     return duals
 end
@@ -110,18 +120,23 @@ function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
                seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
     offset = index - 1
     idxs = Iterators.drop(structural_eachindex(duals, x), offset)
-    if isbitstype(V)
-        for idx in idxs
-            duals[idx] = Dual{T,V,N}(x[idx], seed)
-        end
-    else
-        for idx in idxs
-            if isassigned(x, idx)
+    if supports_fast_scalar_indexing(duals)
+        if isbitstype(V)
+            for idx in idxs
                 duals[idx] = Dual{T,V,N}(x[idx], seed)
-            else
-                Base._unsetindex!(duals, idx)
+            end
+        else
+            for idx in idxs
+                if isassigned(x, idx)
+                    duals[idx] = Dual{T,V,N}(x[idx], seed)
+                else
+                    Base._unsetindex!(duals, idx)
+                end
             end
         end
+    else
+        idxs = collect(idxs)
+        duals[idxs] .= Dual{T,V,N}.(view(x, idxs), Ref(seed))
     end
     return duals
 end
@@ -130,18 +145,23 @@ function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
                seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
     offset = index - 1
     idxs = Iterators.drop(structural_eachindex(duals, x), offset)
-    if isbitstype(V)
-        for (i, idx) in zip(1:chunksize, idxs)
-            duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
-        end
-    else
-        for (i, idx) in zip(1:chunksize, idxs)
-            if isassigned(x, idx)
+    if supports_fast_scalar_indexing(duals)
+        if isbitstype(V)
+            for (i, idx) in zip(1:chunksize, idxs)
                 duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
-            else
-                Base._unsetindex!(duals, idx)
+            end
+        else
+            for (i, idx) in zip(1:chunksize, idxs)
+                if isassigned(x, idx)
+                    duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
+                else
+                    Base._unsetindex!(duals, idx)
+                end
             end
         end
+    else
+        idxs = collect(Iterators.take(idxs, chunksize))
+        duals[idxs] .= Dual{T,V,N}.(view(x, idxs), getindex.(Ref(seeds), 1:length(idxs)))
     end
     return duals
 end

src/apiutils.jl

@@ -65,17 +65,29 @@ end
 extract_gradient!(::Type{T}, result::AbstractArray, y::Real) where {T} = fill!(result, zero(y))
 function extract_gradient!(::Type{T}, result::AbstractArray, dual::Dual) where {T}
     idxs = structural_eachindex(result)
-    for (i, idx) in zip(1:npartials(dual), idxs)
-        result[idx] = partials(T, dual, i)
+    if supports_fast_scalar_indexing(result)
+        for (i, idx) in zip(1:npartials(dual), idxs)
+            result[idx] = partials(T, dual, i)
+        end
+    else
+        fn = PartialsFn{T}(dual)
+        idxs = collect(Iterators.take(idxs, npartials(dual)))
+        result[idxs] .= fn.(1:length(idxs))
+        return result
     end
     return result
 end
 
 function extract_gradient_chunk!(::Type{T}, result, dual, index, chunksize) where {T}
-    offset = index - 1
-    idxs = Iterators.drop(structural_eachindex(result), offset)
-    for (i, idx) in zip(1:chunksize, idxs)
-        result[idx] = partials(T, dual, i)
+    idxs = Iterators.drop(structural_eachindex(result), index - 1)
+    if supports_fast_scalar_indexing(result)
+        for (i, idx) in zip(1:chunksize, idxs)
+            result[idx] = partials(T, dual, i)
+        end
+    else
+        fn = PartialsFn{T}(dual)
+        idxs = collect(Iterators.take(idxs, chunksize))
+        result[idxs] .= fn.(1:length(idxs))
     end
     return result
 end

src/gradient.jl

@@ -0,0 +1,15 @@
+# overload for array types that
+@inline supports_fast_scalar_indexing(::Array) = true
+
+@inline function supports_fast_scalar_indexing(x::AbstractArray)
+    parent(x) === x && return false
+    return supports_fast_scalar_indexing(parent(x))
+end
+
+# Helper function for broadcasting
+struct PartialsFn{T,D<:Dual}
+    dual::D
+end
+PartialsFn{T}(dual::Dual) where {T} = PartialsFn{T,typeof(dual)}(dual)
+
+(f::PartialsFn{T})(i) where {T} = partials(T, f.dual, i)

src/utils.jl

@@ -118,6 +118,7 @@ ForwardDiff.:≺(::Type{OuterTestTag}, ::Type{TestTag}) = false
         @test precision(typeof(FDNUM)) === precision(V)
         @test precision(NESTED_FDNUM) === precision(PRIMAL)
         @test precision(typeof(NESTED_FDNUM)) === precision(V)
+        
         @test precision(FDNUM; base=10) === precision(PRIMAL; base=10)
         @test precision(typeof(FDNUM); base=10) === precision(V; base=10)
         @test precision(NESTED_FDNUM; base=10) === precision(PRIMAL; base=10)

test/DualTest.jl

@@ -277,4 +277,9 @@ end
     @test Array(grad_jl) ≈ grad
 end
 
+@testset "Scalar Indexing Checks" begin
+    @test ForwardDiff.supports_fast_scalar_indexing(UnitLowerTriangular(view(rand(6, 6), 1:3, 1:3)))
+    @test !ForwardDiff.supports_fast_scalar_indexing(UnitLowerTriangular(view(JLArray(rand(6, 6)), 1:3, 1:3)))
+end
+
 end # module