use an extension instead of Requires on 1.9+ (#170)

* move required code to `__init__` function

* move require code to separate files

* use an extension instead of Requires on 1.9+

* make ChainRulesCore usage into an extension on 1.9+

Author: KristofferC

Project.toml

@@ -9,6 +9,16 @@ DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
+
+[extensions]
+TullioCUDAExt = "CUDA"
+TullioFillArraysExt = "FillArrays"
+TullioTrackerExt = "Tracker"
+
 [compat]
 CUDA = "3.6"
 CUDAKernels = "0.3.3, 0.4"

ext/TullioCUDAExt.jl

@@ -0,0 +1,20 @@
+module TullioCUDAExt
+
+if !isdefined(Base, :get_extension)
+    using ..Tullio, ..CUDA
+else
+    using Tullio, CUDA
+end
+
+Tullio.threader(fun!::F, ::Type{T},
+    Z::AbstractArray, As::Tuple, Is::Tuple, Js::Tuple,
+    redfun, block=0, keep=nothing) where {F<:Function, T<:CUDA.CuArray} =
+    fun!(T, Z, As..., Is..., Js..., keep)
+
+Tullio.∇threader(fun!::F, ::Type{T},
+    As::Tuple, Is::Tuple, Js::Tuple, block=0) where {F<:Function, T<:CUDA.CuArray} =
+    fun!(T, As..., Is..., Js...,)
+
+# Tullio.thread_scalar ... ought to work? Was never fast.
+
+end

ext/TullioChainRulesCoreExt.jl

@@ -0,0 +1,20 @@
+module TullioChainRulesCoreExt
+
+if !isdefined(Base, :get_extension)
+    using ..Tullio, ..ChainRulesCore
+else
+    using Tullio, ChainRulesCore
+end
+
+function ChainRulesCore.rrule(ev::Tullio.Eval, args...)
+    Z = ev.fwd(args...)
+    Z, function tullio_back(Δ)
+        isnothing(ev.rev) && error("no gradient definition here!")
+        dxs = map(ev.rev(Δ, Z, args...)) do dx
+            dx === nothing ? ChainRulesCore.ZeroTangent() : dx
+        end
+        tuple(ChainRulesCore.ZeroTangent(), dxs...)
+    end
+end
+
+end

ext/TullioFillArraysExt.jl

@@ -0,0 +1,12 @@
+module TullioFillArraysExt
+
+if !isdefined(Base, :get_extension)
+    using ..Tullio, ..FillArrays
+else
+    using Tullio, FillArrays
+end
+
+Tullio.promote_storage(::Type{T}, ::Type{F}) where {T, F<:FillArrays.Fill} = T
+Tullio.promote_storage(::Type{F}, ::Type{T}) where {T, F<:FillArrays.Fill} = T
+
+end

ext/TullioTrackerExt.jl

@@ -0,0 +1,21 @@
+module TullioTrackerExt
+
+if !isdefined(Base, :get_extension)
+    using ..Tullio, ..Tracker
+else
+    using Tullio, Tracker
+end
+
+(ev::Tullio.Eval)(A::Tracker.TrackedArray, args...) = Tracker.track(ev, A, args...)
+(ev::Tullio.Eval)(A, B::Tracker.TrackedArray, args...) = Tracker.track(ev, A, B, args...)
+(ev::Tullio.Eval)(A::Tracker.TrackedArray, B::Tracker.TrackedArray, args...) = Tracker.track(ev, A, B, args...)
+
+Tracker.@grad function (ev::Tullio.Eval)(args...)
+    Z = ev.fwd(Tracker.data.(args)...)
+    Z, Δ -> begin
+        isnothing(ev.rev) && error("no gradient definition here!")
+        tuple(ev.rev(Tracker.data(Δ), Z, Tracker.data.(args)...)...)
+    end
+end
+
+end

src/eval.jl

@@ -36,14 +36,12 @@ Base.getindex(o::OneBox, i::Integer...) = o.val
 #========== gradient hooks ==========#
 # Macros like @adjoint need to be hidden behind include(), it seems:
 
-using Requires
-
 # @init @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" include("grad/zygote.jl")
 
-@init @require Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" include("grad/tracker.jl")
-
 # @init @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" include("grad/reverse.jl")
 
+# Provided via en extension on 1.9+
+if !isdefined(Base, :get_extension)
 import ChainRulesCore
 
 function ChainRulesCore.rrule(ev::Eval, args...)
@@ -56,11 +54,18 @@ function ChainRulesCore.rrule(ev::Eval, args...)
         tuple(ChainRulesCore.ZeroTangent(), dxs...)
     end
 end
+end
 
-@init @require FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" begin
-    using .FillArrays: Fill # used by Zygote
-    Tullio.promote_storage(::Type{T}, ::Type{F}) where {T, F<:Fill} = T
-    Tullio.promote_storage(::Type{F}, ::Type{T}) where {T, F<:Fill} = T
+if !isdefined(Base, :get_extension)
+using Requires
+end
+
+@static if !isdefined(Base, :get_extension)
+function __init__()
+    @require Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" include("../ext/TullioTrackerExt.jl")
+    @require FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" include("../ext/TullioFillArraysExt.jl")
+    @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" include("../ext/TullioCUDAExt.jl")
+end
 end
 
 #========== vectorised gradients ==========#
@@ -92,25 +97,6 @@ end
 
 =#
 
-#========== CuArrays ==========#
-
-using Requires
-
-@init @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" begin
-    using .CUDA: CuArray, GPUArrays
-
-    Tullio.threader(fun!::F, ::Type{T},
-        Z::AbstractArray, As::Tuple, Is::Tuple, Js::Tuple,
-        redfun, block=0, keep=nothing) where {F<:Function, T<:CuArray} =
-        fun!(T, Z, As..., Is..., Js..., keep)
-
-    Tullio.∇threader(fun!::F, ::Type{T},
-        As::Tuple, Is::Tuple, Js::Tuple, block=0) where {F<:Function, T<:CuArray} =
-        fun!(T, As..., Is..., Js...,)
-
-    # Tullio.thread_scalar ... ought to work? Was never fast.
-end
-
 #========== storage unwrapper ==========#
 
 """