Merge pull request #31 from ziyiyin97/ad

wrap a function as a jet operator

Author: samtkaplan

Project.toml

@@ -1,16 +1,27 @@
 name = "JetPack"
 uuid = "24ef3835-3876-54c3-8a7a-956cf69ca0b2"
-version = "1.1.2"
+version = "1.2.0"
 
 [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"
 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"
 SpecialFunctions = "1, 2"
 julia = "1"
+
+[extras]
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"

README.md

@@ -13,6 +13,7 @@ These operators include the following functionality:
 * Blending and projection operators
 * Operators implementing transcendental functions like logarithm and exponentiation, with linearization
 * Special functions as used in some full waveform inversion approaches, including *Normalized Integral Method* and linear non-analytic operators including taking real and imaginary parts
+* Wrapping an auto-differentiable Julia function as a Jet operator
 
 [docs-dev-img]: https://img.shields.io/badge/docs-dev-blue.svg
 [docs-dev-url]: https://chevronetc.github.io/JetPack.jl/dev/

ext/JetPackFluxExt.jl

@@ -0,0 +1,21 @@
+module JetPackFluxExt
+
+using JetPack, Jets, Flux
+"""
+    F = JopAD(f, dom, rng)
+
+where `F` wraps a Julia native function `f` as a Jet operator with domain
+and range given by `dom` and `rng`, respectively. Notice that it is the
+responsibility of the user to ensure that function `f` is auto-differentiable
+in both forward-mode and reverse-mode via the function `pushforward` and
+`pullback` in Flux. If the user chooses to use customized AD rules for this
+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).
+""" 
+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]
+
+end

src/JetPack.jl

@@ -42,4 +42,14 @@ include("jop_taper.jl")
 include("jop_tanh.jl")
 include("jop_translation.jl")
 
+
+###### JetPack with Flux extension
+
+function JopAD end;
+export JopAD
+
+if !isdefined(Base, :get_extension)
+    include("../ext/JetPackFluxExt.jl")
+end
+
 end

test/Project.toml

@@ -1,5 +1,6 @@
 [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"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

test/jop_ad.jl

@@ -0,0 +1,24 @@
+using LinearAlgebra, Jets, JetPack, Test, Flux
+
+spc = JetSpace(Float64, 10, 10)
+
+### set up a nonlinear operator according to the function `f`
+f(x) = x ^ 2
+F = JopAD(f; dom=spc, rng=spc)
+
+@testset "JopAD: linearity, dot product, and linearization tests" begin
+    
+    x = rand(domain(F))
+    J = jacobian(F, x)
+    lhs, rhs = dot_product_test(J, rand(domain(J)), rand(range(J)))
+    @test lhs ≈ rhs
+    lhs,rhs = linearity_test(J)
+    @test lhs ≈ rhs
+    m0 = rand(domain(F))
+    μ  = sqrt.([1/1,1/2,1/4,1/8,1/16,1/32,1/64,1/128,1/256,1/512,1/1024,1/2048,1/4096,1/8192])
+    δm = rand(domain(F))
+    observed, expected = linearization_test(F, m0, μ = μ, δm = δm)
+    δ = minimum(abs, observed - expected)
+    @test δ < 1e-6
+    
+end

test/runtests.jl

@@ -4,6 +4,7 @@ Random.seed!(101)
 
 for filename in (
         "jop_atan.jl",
+        "jop_ad.jl",
         "jop_blend.jl",
         "jop_circshift.jl",
         "jop_derivative.jl",