fix: support gradient + more test coverage

Author: avik-pal

Project.toml

@@ -6,6 +6,7 @@ version = "1.0.2"
 CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"

ext/ForwardDiffGPUArraysCoreExt.jl

@@ -1,36 +1,62 @@
 module ForwardDiffGPUArraysCoreExt
 
 using GPUArraysCore: AbstractGPUArray
-using ForwardDiff: ForwardDiff, Dual, Partials
+using ForwardDiff: ForwardDiff, Dual, Partials, npartials, partials
 
 function ForwardDiff.seed!(duals::AbstractGPUArray{Dual{T,V,N}}, x,
                            seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
-    idxs = ForwardDiff.structural_eachindex(duals, x)
-    duals[idxs] .= Dual{T,V,N}.(x[idxs], seed)
+    idxs = collect(ForwardDiff.structural_eachindex(duals, x))
+    duals[idxs] .= Dual{T,V,N}.(x[idxs], Ref(seed))
     return duals
 end
 
 function ForwardDiff.seed!(duals::AbstractGPUArray{Dual{T,V,N}}, x,
                            seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
-    idxs = ForwardDiff.structural_eachindex(duals, x)
+    idxs = collect(ForwardDiff.structural_eachindex(duals, x))[1:N]
     duals[idxs] .= Dual{T,V,N}.(x[idxs], seeds[1:N])
     return duals
 end
 
 function ForwardDiff.seed!(duals::AbstractGPUArray{Dual{T,V,N}}, x, index,
                            seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
     offset = index - 1
-    idxs = Iterators.drop(ForwardDiff.structural_eachindex(duals, x), offset)
-    duals[idxs] .= Dual{T,V,N}.(x[idxs], seed)
+    idxs = collect(Iterators.drop(ForwardDiff.structural_eachindex(duals, x), offset))
+    duals[idxs] .= Dual{T,V,N}.(x[idxs], Ref(seed))
     return duals
 end
 
 function ForwardDiff.seed!(duals::AbstractGPUArray{Dual{T,V,N}}, x, index,
                            seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
     offset = index - 1
-    idxs = Iterators.drop(ForwardDiff.structural_eachindex(duals, x), offset)
+    idxs = collect(
+        Iterators.drop(ForwardDiff.structural_eachindex(duals, x), offset)
+    )[1:chunksize]
     duals[idxs] .= Dual{T,V,N}.(x[idxs], seeds[1:chunksize])
     return duals
 end
 
+# gradient
+function ForwardDiff.extract_gradient!(::Type{T}, result::AbstractGPUArray,
+                                       dual::Dual) where {T}
+    # this closure is needed for gpu compilation
+    partial_fn(dual, i) = partials(T, dual, i)
+
+    idxs = ForwardDiff.structural_eachindex(result)
+    result[idxs] .= partial_fn.(Ref(dual), 1:npartials(dual))
+    return result
+end
+
+function ForwardDiff.extract_gradient_chunk!(::Type{T}, result::AbstractGPUArray, dual,
+                                             index, chunksize) where {T}
+    # this closure is needed for gpu compilation
+    partial_fn(dual, i) = partials(T, dual, i)
+
+    offset = index - 1
+    idxs = collect(
+        Iterators.drop(ForwardDiff.structural_eachindex(result), offset)
+    )[1:chunksize]
+    result[idxs] .= partial_fn.(Ref(dual), 1:chunksize)
+    return result
+end
+
 end

test/GradientTest.jl

@@ -9,6 +9,7 @@ using ForwardDiff
 using ForwardDiff: Dual, Tag
 using StaticArrays
 using DiffTests
+using JLArrays
 
 include(joinpath(dirname(@__FILE__), "utils.jl"))
 
@@ -255,4 +256,25 @@ end
     end
 end
 
+@testset "GPUArraysCore" begin
+    fn(x) = sum(x .^ 2 ./ 2)
+
+    x = [1.0, 2.0, 3.0]
+    x_jl = JLArray(x)
+
+    grad = ForwardDiff.gradient(fn, x)
+    grad_jl = ForwardDiff.gradient(fn, x_jl)
+
+    @test grad_jl isa JLArray
+    @test Array(grad_jl) ≈ grad
+
+    cfg = ForwardDiff.GradientConfig(
+        fn, x_jl, ForwardDiff.Chunk{2}(), ForwardDiff.Tag(fn, eltype(x))
+    )
+    grad_jl = ForwardDiff.gradient(fn, x_jl, cfg)
+
+    @test grad_jl isa JLArray
+    @test Array(grad_jl) ≈ grad
+end
+
 end # module