Merge pull request #22 from LuxDL/ap/simplify_ext

Use PackageExtensionCompat
LuxDL · Jun 26, 2023 · 6980015 · 6980015 · avik-pal · Jun 26, 2023
2 parents d3274eb + a708262
commit 6980015
Show file tree

Hide file tree

Showing 7 changed files with 22 additions and 82 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,16 +1,16 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.2.4"
+version = "0.2.5"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [weakdeps]
@@ -32,8 +32,8 @@ ForwardDiff = "0.10"
 KernelAbstractions = "0.9"
 LuxCUDA = "0.1"
 NNlib = "0.8, 0.9"
+PackageExtensionCompat = "1"
 Reexport = "1"
-Requires = "1"
 ReverseDiff = "1"
 Tracker = "0.2"
 julia = "1.6"

diff --git a/ext/LuxLibForwardDiffExt.jl b/ext/LuxLibForwardDiffExt.jl
@@ -1,7 +1,6 @@
 module LuxLibForwardDiffExt
 
-isdefined(Base, :get_extension) ? (using ForwardDiff) : (using ..ForwardDiff)
-using LuxLib
+using ForwardDiff, LuxLib
 
 function LuxLib._dropout_fptype(x::AbstractArray{<:ForwardDiff.Dual})
     return ForwardDiff.valtype(eltype(x))

diff --git a/ext/LuxLibLuxCUDAExt.jl b/ext/LuxLibLuxCUDAExt.jl
@@ -1,7 +1,6 @@
 module LuxLibLuxCUDAExt
 
-isdefined(Base, :get_extension) ? (using LuxCUDA) : (using ..LuxCUDA)
-using LuxLib
+using LuxCUDA, LuxLib
 import ChainRulesCore as CRC
 import LuxLib: batchnorm, _batchnorm_cudnn!, _get_batchnorm_statistics, FP_32_64, ∂∅
 

diff --git a/ext/LuxLibLuxCUDATrackerExt.jl b/ext/LuxLibLuxCUDATrackerExt.jl
@@ -1,16 +1,7 @@
 module LuxLibLuxCUDATrackerExt
 
-if isdefined(Base, :get_extension)
-    using Tracker
-    import Tracker: @grad, data, nobacksies, track, TrackedArray, TrackedVector, TrackedReal
-    using LuxCUDA
-else
-    using ..Tracker
-    import ..Tracker: @grad,
-        data, nobacksies, track, TrackedArray, TrackedVector, TrackedReal
-    using ..LuxCUDA
-end
-using NNlib, LuxLib
+using NNlib, LuxCUDA, LuxLib, Tracker
+import Tracker: @grad, data, nobacksies, track, TrackedArray, TrackedVector, TrackedReal
 import LuxLib: AA,
     AV, _batchnorm_cudnn!, _get_batchnorm_statistics, FP_32_64, ∂∅, __is_tracked
 

diff --git a/ext/LuxLibReverseDiffExt.jl b/ext/LuxLibReverseDiffExt.jl
@@ -1,33 +1,18 @@
 module LuxLibReverseDiffExt
 
-if isdefined(Base, :get_extension)
-    using ReverseDiff
-    import ReverseDiff: SpecialInstruction,
-        TrackedArray,
-        TrackedReal,
-        decrement_deriv!,
-        increment_deriv!,
-        track,
-        value,
-        special_reverse_exec!,
-        special_forward_exec!,
-        @grad_from_chainrules
-else
-    using ..ReverseDiff
-    import ..ReverseDiff: SpecialInstruction,
-        TrackedArray,
-        TrackedReal,
-        decrement_deriv!,
-        increment_deriv!,
-        track,
-        value,
-        special_reverse_exec!,
-        special_forward_exec!,
-        @grad_from_chainrules
-end
-using ChainRulesCore, LuxLib, NNlib
+using ChainRulesCore, LuxLib, NNlib, ReverseDiff
 import ChainRulesCore as CRC
 import LuxLib: AA, __is_tracked
+import ReverseDiff: SpecialInstruction,
+    TrackedArray,
+    TrackedReal,
+    decrement_deriv!,
+    increment_deriv!,
+    track,
+    value,
+    special_reverse_exec!,
+    special_forward_exec!,
+    @grad_from_chainrules
 
 # Patches: Needs upstreaming
 @inline function increment_deriv!(t::Union{TrackedArray, TrackedReal}, ::NoTangent, i)

diff --git a/ext/LuxLibTrackerExt.jl b/ext/LuxLibTrackerExt.jl
@@ -1,17 +1,10 @@
 module LuxLibTrackerExt
 
-if isdefined(Base, :get_extension)
-    using Tracker
-    import Tracker: @grad, data, nobacksies, track, TrackedArray, TrackedVector, TrackedReal
-else
-    using ..Tracker
-    import ..Tracker: @grad,
-        data, nobacksies, track, TrackedArray, TrackedVector, TrackedReal
-end
-using NNlib, LuxLib
+using NNlib, LuxLib, Tracker
 import LuxLib: AA,
     AV, _batchnorm_cudnn!, _get_batchnorm_statistics, FP_32_64, ∂∅, __is_tracked
 import ChainRulesCore as CRC
+import Tracker: @grad, data, nobacksies, track, TrackedArray, TrackedVector, TrackedReal
 
 # NNlib: batched_mul
 for T1 in (:AbstractArray, :TrackedArray), T2 in (:AbstractArray, :TrackedArray)

diff --git a/src/LuxLib.jl b/src/LuxLib.jl
@@ -11,36 +11,9 @@ using KernelAbstractions
 import KernelAbstractions as KA
 
 # Extensions
-if !isdefined(Base, :get_extension)
-    using Requires
-end
-
+using PackageExtensionCompat
 function __init__()
-    @static if !isdefined(Base, :get_extension)
-        # Handling AD Packages
-        ## Handling ForwardDiff
-        @require ForwardDiff="f6369f11-7733-5829-9624-2563aa707210" begin
-            include("../ext/LuxLibForwardDiffExt.jl")
-        end
-        ## Handling Tracker
-        @require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" begin
-            include("../ext/LuxLibTrackerExt.jl")
-        end
-        ## Handling ReverseDiff
-        @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
-            include("../ext/LuxLibReverseDiffExt.jl")
-        end
-
-        # Accelerator Support
-        ## Handling CUDA
-        @require LuxCUDA="d0bbae9a-e099-4d5b-a835-1c6931763bda" begin
-            include("../ext/LuxLibLuxCUDAExt.jl")
-
-            @require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" begin
-                include("../ext/LuxLibLuxCUDATrackerExt.jl")
-            end
-        end
-    end
+    @require_extensions
 end
 
 include("utils.jl")