fix: apply suggestions from code review

Co-authored-by: David Widmann <[email protected]>

Author: avik-pal

ext/ForwardDiffGPUArraysCoreExt.jl

@@ -4,34 +4,37 @@ using GPUArraysCore: AbstractGPUArray
 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}
+                           seed::Partials{N,V}) where {T,V,N}
     idxs = collect(ForwardDiff.structural_eachindex(duals, x))
-    duals[idxs] .= Dual{T,V,N}.(x[idxs], Ref(seed))
+    duals[idxs] .= Dual{T,V,N}.(view(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 = collect(ForwardDiff.structural_eachindex(duals, x))[1:N]
-    duals[idxs] .= Dual{T,V,N}.(x[idxs], seeds[1:N])
+    idxs = collect(Iterators.take(ForwardDiff.structural_eachindex(duals, x), N))
+    duals[idxs] .= Dual{T,V,N}.(view(x, idxs), getindex.(Ref(seeds), 1:length(idxs)))
     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}
+                           seed::Partials{N,V}) where {T,V,N}
     offset = index - 1
     idxs = collect(Iterators.drop(ForwardDiff.structural_eachindex(duals, x), offset))
-    duals[idxs] .= Dual{T,V,N}.(x[idxs], Ref(seed))
+    duals[idxs] .= Dual{T,V,N}.(view(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}
+                           seeds::NTuple{N,Partials{N,V}}, chunksize) where {T,V,N}
     offset = index - 1
     idxs = collect(
-        Iterators.drop(ForwardDiff.structural_eachindex(duals, x), offset)
-    )[1:chunksize]
-    duals[idxs] .= Dual{T,V,N}.(x[idxs], seeds[1:chunksize])
+        Iterators.take(
+            Iterators.drop(ForwardDiff.structural_eachindex(duals, x), offset),
+            chunksize,
+        ),
+    )
+    duals[idxs] .= Dual{T,V,N}.(view(x, idxs), getindex.(Ref(seeds), 1:length(idxs)))
     return duals
 end
 
@@ -41,8 +44,8 @@ function ForwardDiff.extract_gradient!(::Type{T}, result::AbstractGPUArray,
     # 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))
+    idxs = collect(Iterators.take(ForwardDiff.structural_eachindex(result), npartials(dual)))
+    result[idxs] .= partial_fn.(Ref(dual), 1:length(idxs))
     return result
 end
 
@@ -53,9 +56,12 @@ function ForwardDiff.extract_gradient_chunk!(::Type{T}, result::AbstractGPUArray
 
     offset = index - 1
     idxs = collect(
-        Iterators.drop(ForwardDiff.structural_eachindex(result), offset)
-    )[1:chunksize]
-    result[idxs] .= partial_fn.(Ref(dual), 1:chunksize)
+        Iterators.take(
+                    Iterators.drop(ForwardDiff.structural_eachindex(result), offset),
+            chunksize,
+        )
+    )
+    result[idxs] .= partial_fn.(Ref(dual), 1:length(idxs))
     return result
 end