refactor: move ForwardDiff to an ext

Author: avik-pal

Project.toml

@@ -19,6 +19,7 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [extensions]
@@ -27,6 +28,7 @@ NNlibCUDACUDNNExt = ["CUDA", "cuDNN"]
 NNlibCUDAExt = "CUDA"
 NNlibEnzymeCoreExt = "EnzymeCore"
 NNlibFFTWExt = "FFTW"
+NNlibForwardDiffExt = "ForwardDiff"
 
 [compat]
 AMDGPU = "0.9.4, 1"
@@ -36,6 +38,7 @@ CUDA = "4, 5"
 ChainRulesCore = "1.13"
 EnzymeCore = "0.5, 0.6, 0.7"
 FFTW = "1.8.0"
+ForwardDiff = "0.10.36"
 GPUArraysCore = "0.1"
 KernelAbstractions = "0.9.2"
 LinearAlgebra = "<0.0.1, 1"

ext/NNlibForwardDiffExt.jl

@@ -0,0 +1,9 @@
+module NNlibForwardDiffExt
+
+using ForwardDiff: ForwardDiff
+using NNlib: NNlib
+
+NNlib.within_gradient(x::ForwardDiff.Dual) = true
+NNlib.within_gradient(x::AbstractArray{<:ForwardDiff.Dual}) = true
+
+end

src/NNlib.jl

@@ -82,11 +82,6 @@ export upsample_nearest, ∇upsample_nearest,
 include("gather.jl")
 include("scatter.jl")
 include("utils.jl")
-@init @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" begin
-    using .ForwardDiff
-    within_gradient(x::ForwardDiff.Dual) = true
-    within_gradient(x::AbstractArray{<:ForwardDiff.Dual}) = true
-end
 
 include("sampling.jl")
 include("functions.jl")