fix: forwarddiff support for gpu arrays (#1605)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.28.0"
+version = "1.29.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -101,7 +101,7 @@ Enzyme = "0.13.81"
 EnzymeCore = "0.8.15"
 FastClosures = "0.3.2"
 Flux = "0.16.3"
-ForwardDiff = "0.10.36, =1"
+ForwardDiff = "0.10.36, 1"
 FunctionWrappers = "1.1.3"
 Functors = "0.5"
 GPUArrays = "11"

ext/LuxComponentArraysExt.jl

@@ -2,6 +2,7 @@ module LuxComponentArraysExt
 
 using ComponentArrays: ComponentArrays, ComponentArray
 using Lux: Lux, DistributedUtils
+using ForwardDiff: ForwardDiff
 
 # Distributed Functionality
 function DistributedUtils.synchronize!!(
@@ -11,4 +12,9 @@ function DistributedUtils.synchronize!!(
     return ComponentArray(ps_synced, ComponentArrays.getaxes(ps))
 end
 
+@static if pkgversion(ForwardDiff) ≥ v"1.0.1"
+    # Apply overloads for GPU arrays
+    Lux.@define_forwarddiff_gpu_overloads ComponentArray
+end
+
 end

src/helpers/forwarddiff_training.jl

@@ -89,3 +89,113 @@ function Training.compute_gradients_impl(
         ts,
     )
 end
+
+# Type Piracy for ForwardDiff GPU Array Support
+# This is a workaround for ForwardDiff.jl not supporting GPU arrays post v1.0
+# See: https://github.com/JuliaDiff/ForwardDiff.jl/pull/760
+
+using GPUArraysCore: AnyGPUArray
+
+# Helper struct for broadcasting partials extraction
+struct PartialsFn{T,D<:ForwardDiff.Dual}
+    dual::D
+end
+
+PartialsFn{T}(dual::ForwardDiff.Dual) where {T} = PartialsFn{T,typeof(dual)}(dual)
+
+(f::PartialsFn{T})(i) where {T} = ForwardDiff.partials(T, f.dual, i)
+
+# Macro to define ForwardDiff overloads for array types that don't support scalar indexing
+macro define_forwarddiff_gpu_overloads(ArrayType)
+    return quote
+        # Overloaded seed! methods
+        function ForwardDiff.seed!(
+            duals::$(esc(ArrayType)){ForwardDiff.Dual{T,V,N}},
+            x,
+            seed::ForwardDiff.Partials{N,V}=zero(ForwardDiff.Partials{N,V}),
+        ) where {T,V,N}
+            idxs = collect(ForwardDiff.structural_eachindex(duals, x))
+            duals[idxs] .= ForwardDiff.Dual{T,V,N}.(view(x, idxs), Ref(seed))
+            return duals
+        end
+
+        function ForwardDiff.seed!(
+            duals::$(esc(ArrayType)){ForwardDiff.Dual{T,V,N}},
+            x,
+            seeds::NTuple{N,ForwardDiff.Partials{N,V}},
+        ) where {T,V,N}
+            idxs = collect(Iterators.take(ForwardDiff.structural_eachindex(duals, x), N))
+            duals[idxs] .=
+                ForwardDiff.Dual{
+                    T,V,N
+                }.(view(x, idxs), getindex.(Ref(seeds), 1:length(idxs)))
+            return duals
+        end
+
+        function ForwardDiff.seed!(
+            duals::$(esc(ArrayType)){ForwardDiff.Dual{T,V,N}},
+            x,
+            index,
+            seed::ForwardDiff.Partials{N,V}=zero(ForwardDiff.Partials{N,V}),
+        ) where {T,V,N}
+            idxs = collect(
+                Iterators.drop(ForwardDiff.structural_eachindex(duals, x), index - 1)
+            )
+            duals[idxs] .= ForwardDiff.Dual{T,V,N}.(view(x, idxs), Ref(seed))
+            return duals
+        end
+
+        function ForwardDiff.seed!(
+            duals::$(esc(ArrayType)){ForwardDiff.Dual{T,V,N}},
+            x,
+            index,
+            seeds::NTuple{N,ForwardDiff.Partials{N,V}},
+            chunksize=N,
+        ) where {T,V,N}
+            idxs = collect(
+                Iterators.take(
+                    Iterators.drop(ForwardDiff.structural_eachindex(duals, x), index - 1),
+                    chunksize,
+                ),
+            )
+            duals[idxs] .=
+                ForwardDiff.Dual{
+                    T,V,N
+                }.(view(x, idxs), getindex.(Ref(seeds), 1:length(idxs)))
+            return duals
+        end
+
+        # Overloaded extract_gradient! methods
+        function ForwardDiff.extract_gradient!(
+            ::Type{T}, result::$(esc(ArrayType)), dual::ForwardDiff.Dual
+        ) where {T}
+            fn = PartialsFn{T}(dual)
+            idxs = collect(
+                Iterators.take(
+                    ForwardDiff.structural_eachindex(result), ForwardDiff.npartials(dual)
+                ),
+            )
+            result[idxs] .= fn.(1:length(idxs))
+            return result
+        end
+
+        function ForwardDiff.extract_gradient_chunk!(
+            ::Type{T}, result::$(esc(ArrayType)), dual, index, chunksize
+        ) where {T}
+            fn = PartialsFn{T}(dual)
+            idxs = collect(
+                Iterators.take(
+                    Iterators.drop(ForwardDiff.structural_eachindex(result), index - 1),
+                    chunksize,
+                ),
+            )
+            result[idxs] .= fn.(1:length(idxs))
+            return result
+        end
+    end
+end
+
+@static if pkgversion(ForwardDiff) ≥ v"1.0.1"
+    # Apply overloads for GPU arrays
+    @define_forwarddiff_gpu_overloads AnyGPUArray
+end

test/qa_tests.jl

@@ -9,12 +9,21 @@
         exclude=[
             ForwardDiff.jacobian,
             ForwardDiff.gradient,
+            ForwardDiff.extract_gradient_chunk!,
             Lux.AutoDiffInternalImpl.batched_jacobian,
             Lux.AutoDiffInternalImpl.jacobian_vector_product,
             Lux.AutoDiffInternalImpl.jacobian_vector_product_impl,
         ],
     )
-    Aqua.test_piracies(Lux; treat_as_own=[Lux.outputsize])
+    Aqua.test_piracies(
+        Lux;
+        treat_as_own=[
+            Lux.outputsize,
+            ForwardDiff.extract_gradient_chunk!,
+            ForwardDiff.extract_gradient!,
+            ForwardDiff.seed!,
+        ],
+    )
     Aqua.test_unbound_args(Lux; broken=true)
 end