Automatically register GPU converters upon loading the glue package.

Author: maleadt

lib/FluxAMDGPU/src/FluxAMDGPU.jl

@@ -10,4 +10,13 @@ using Flux: OneHotArray, OneHotLike
 Base.BroadcastStyle(::Type{<:OneHotArray{<:Any, <:Any, <:Any, N, <:ROCArray}}) where N =
     AMDGPU.ROCArrayStyle{N}()
 
+function __init__()
+    if Flux.default_gpu_converter[] === identity
+        @info "Registering AMDGPU.jl as the default GPU converter"
+        Flux.default_gpu_converter[] = roc
+    else
+        @warn "Not registering AMDGPU.jl as the default GPU converter as another one has been registered already."
+    end
+end
+
 end # module

lib/FluxAMDGPU/test/runtests.jl

@@ -9,7 +9,7 @@ using .AMDGPU
 ENV["JULIA_GPU_ALLOWSCALAR"] = "false"
 
 using .Flux
-Flux.default_gpu_converter[] = AMDGPU.roc
+@assert Flux.default_gpu_converter[] == roc
 
 using Zygote
 using Zygote: pullback

lib/FluxCUDA/src/FluxCUDA.jl

@@ -9,5 +9,13 @@ include("onehot.jl")
 include("ctc.jl")
 include("cudnn.jl")
 
+function __init__()
+    if Flux.default_gpu_converter[] === identity
+        @info "Registering CUDA.jl as the default GPU converter"
+        Flux.default_gpu_converter[] = cu
+    else
+        @warn "Not registering CUDA.jl as the default GPU converter as another one has been registered already."
+    end
+end
 
 end # module

lib/FluxCUDA/test/runtests.jl

@@ -9,7 +9,7 @@ using .CUDA
 ENV["JULIA_GPU_ALLOWSCALAR"] = "false"
 
 using .Flux
-Flux.default_gpu_converter[] = cu
+@assert Flux.default_gpu_converter[] == cu
 
 using Zygote
 using Zygote: pullback