switch from Requires.jl to an extension

Author: ziyiyin97

Project.toml

@@ -4,15 +4,24 @@ version = "1.1.3"
 
 [deps]
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Jets = "2a57b368-ab28-5ba9-84aa-637fe7991822"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
+[weakdeps]
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+
+[extensions]
+JetPackFluxExt = "Flux"
+
 [compat]
 FFTW = "1"
+Flux = "0.12, 0.13"
 Jets = "^1.2"
-Requires = "1"
 SpecialFunctions = "1, 2"
 julia = "1"
+
+[extras]
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"

src/jop_ad.jl ext/JetPackFluxExt.jlsrc/jop_ad.jl renamed to ext/JetPackFluxExt.jl

@@ -1,3 +1,6 @@
+module JetPackFluxExt
+
+using JetPack, Jets, Flux
 """
     F = JopAD(f, dom, rng)
 
@@ -9,9 +12,10 @@ in both forward-mode and reverse-mode via the function `pushforward` and
 function `f`, it is also the responsibility of the user to ensure that these
 rules are tested (e.g. linearity, dot product, and linearization tests).
 """ 
-JopAD(f::Function; dom, rng) = JopNl(dom = dom, rng = rng, f! = JopAD_f!, df! = JopAD_df!, df′! = JopAD_df′!, s = (f=f,))
-export JopAD
+JetPack.JopAD(f::Function; dom, rng) = JopNl(dom = dom, rng = rng, f! = JopAD_f!, df! = JopAD_df!, df′! = JopAD_df′!, s = (f=f,))
 
 JopAD_f!(d::AbstractArray, m::AbstractArray; f::Function) = d .= f(m)
 JopAD_df!(δd::AbstractArray, δm::AbstractArray; mₒ, f::Function) = δd .= Flux.pushforward(f, mₒ)(δm)
-JopAD_df′!(δm::AbstractArray, δd::AbstractArray; mₒ, f::Function) = δm .= Flux.pullback(f, mₒ)[2](δd)[1]
+JopAD_df′!(δm::AbstractArray, δd::AbstractArray; mₒ, f::Function) = δm .= Flux.pullback(f, mₒ)[2](δd)[1]
+
+end

src/JetPack.jl

@@ -5,7 +5,6 @@ using Jets
 using LinearAlgebra
 using SpecialFunctions
 using Statistics
-using Requires
 
 include("jop_atan.jl")
 include("jop_blend.jl")
@@ -43,9 +42,7 @@ include("jop_taper.jl")
 include("jop_tanh.jl")
 include("jop_translation.jl")
 
-### If Flux is loaded, users can wrap an auto-differentiable function as a Jet operator
-function __init__()
-    @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" @eval include("jop_ad.jl")
-end
+function JopAD end;
+export JopAD
 
 end