Merge pull request #3 from LuxDL/ap/simplify_ext

Use PackageExtensionCompat

Author: avik-pal

.buildkite/pipeline.yml

@@ -89,4 +89,4 @@ steps:
           soft_fail: true
 
 env:
-  SECRET_CODECOV_TOKEN: "XiQca3XDkJesuEeTkH5zFOrX0zmyXN03NkySFjZFeC37wDqmA6vHlbhDa3XOA4T8b6cNvo4boO72gXlnVkZyPRHVFWPOr338fxAi6Eif7k5TuN44pl2A+DoNZYqM1XyxW8+BR1+zgh1U7wf3PadN5eTtWlZsXUy1ULH8DPaPgqenv9McU3VjsGtaRWQlYplOKZNuVo5HMIdliwWK7eb0ij7QBB4QZNoVAMonXtGE3Q9X2rqMxRky5QmkuaC0RWOdMCAoPe13pj/c1GYSNHXugGiUFDzgyjX/IsK07N+ApzKkqHFp4LEPddhQCD+KU+seMnxl9DHiAOejnrbs1oVXiw==;U2FsdGVkX1/+LzYYK1HvRFpGBhtRqBz4QcrLLtwM2aoMZBDwHsz0VSO3RN4aciB988iEP2xLn24LFtZ4wNS1xg=="
+  SECRET_CODECOV_TOKEN: "PxSr3Y7vdbiwaoX51uGykPsogxmP1IOBt5Z8TwP9GqDxIrvFocEVV2DR4Bebee12G/HYvXtQTyYXH49DpzlsfJ7ri1GQZxd9WRr+aM1DDYmzfDCfpadp4hMoJ5NQvmc/PzeGrNWOOaewaLTUP1eEaG4suygZN0lc5q9BCchIJeqoklGms5DVt/HtfTmwoD/s4wGoIJINi4RoFgnCAkzSh11hTAkyjVerfBGWEi/8E6+WBq3UKwaW4HnT02wG9qFnD4XkHpIpjMxJTpdBn5ufKI+QoJ7qJHlwqgDCtsOCblApccLTjH/BnTahNoSb/b0wdS/cblOTrtdPGzZ5UvmQ4Q==;U2FsdGVkX1/Ji2Nqeq3tqTYCBik6iXILP+rriPRqj/qxhFu4vBWWT3UnlfqDzj6oVdXyuKt0+5e+x33x2S0mBw=="

Project.toml

@@ -8,9 +8,9 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
+PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
@@ -36,8 +36,8 @@ LuxAMDGPU = "0.1"
 LuxCUDA = "0.1"
 LuxCore = "0.1.4"
 Metal = "0.4"
+PackageExtensionCompat = "1"
 Preferences = "1"
-Requires = "1"
 Zygote = "0.6"
 julia = "1.6"
 

ext/LuxDeviceUtilsFillArraysExt.jl

@@ -1,8 +1,6 @@
 module LuxDeviceUtilsFillArraysExt
 
-isdefined(Base, :get_extension) ? (using FillArrays) : (using ..FillArrays)
-
-using Adapt, LuxDeviceUtils
+using Adapt, FillArrays, LuxDeviceUtils
 
 Adapt.adapt_structure(::LuxCPUAdaptor, x::FillArrays.AbstractFill) = x
 

ext/LuxDeviceUtilsLuxCUDAExt.jl

@@ -1,7 +1,6 @@
 module LuxDeviceUtilsLuxCUDAExt
 
-isdefined(Base, :get_extension) ? (using LuxCUDA) : (using ..LuxCUDA)
-using ChainRulesCore, LuxDeviceUtils, Random
+using ChainRulesCore, LuxCUDA, LuxDeviceUtils, Random
 import Adapt: adapt_storage, adapt
 import ChainRulesCore as CRC
 

ext/LuxDeviceUtilsZygoteExt.jl

@@ -1,8 +1,6 @@
 module LuxDeviceUtilsZygoteExt
 
-isdefined(Base, :get_extension) ? (using Zygote) : (using ..Zygote)
-
-using Adapt, LuxDeviceUtils
+using Adapt, LuxDeviceUtils, Zygote
 
 Adapt.adapt_structure(::LuxCPUAdaptor, x::Zygote.OneElement) = x
 

src/LuxDeviceUtils.jl

@@ -4,33 +4,11 @@ using Functors, LuxCore, Preferences, Random, SparseArrays
 import Adapt: adapt, adapt_storage
 import Base: PkgId, UUID
 
-## -----------
-## Extensions
-if !isdefined(Base, :get_extension)
-    using Requires
-end
-
+using PackageExtensionCompat
 function __init__()
-    @static if !isdefined(Base, :get_extension)
-        @require FillArrays="1a297f60-69ca-5386-bcde-b61e274b549b" begin
-            include("../ext/LuxDeviceUtilsFillArraysExt.jl")
-        end
-
-        @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" begin
-            include("../ext/LuxDeviceUtilsZygoteExt.jl")
-        end
-
-        # Accelerators: CUDA Support
-        @require LuxCUDA="d0bbae9a-e099-4d5b-a835-1c6931763bda" begin
-            include("../ext/LuxDeviceUtilsLuxCUDAExt.jl")
-        end
-
-        # NOTE: AMDGPU & Metal Support is only available on Julia 1.9+
-    end
+    @require_extensions
 end
 
-## -----------
-
 export gpu_backend!, supported_gpu_backends
 export gpu_device, cpu_device, LuxCPUDevice, LuxCUDADevice, LuxAMDGPUDevice, LuxMetalDevice
 export LuxCPUAdaptor, LuxCUDAAdaptor, LuxAMDGPUAdaptor, LuxMetalAdaptor