Support AD via Zygote (WIP).

.travis.yml

@@ -20,8 +20,8 @@ matrix:
 #  apt: # apt-get for linux
 #    packages:
 #    - gfortran
-#before_script: # homebrew for mac
-#  - if [ $TRAVIS_OS_NAME = osx ]; then brew install gcc; fi
+before_script: # homebrew for mac
+  - julia -e 'using Pkg; pkg"add Zygote#master IRTools#master"'
 after_success:
   - julia submit_coverage.jl    # submit coverage
 jobs:

Project.toml

@@ -8,10 +8,12 @@ ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+IRTools = "7869d1d1-7146-5819-86e3-90919afe41df"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 TransformVariables = "84d833dd-6860-57f9-a1a7-6da5db126cff"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 TransformVariables = "≥ 0.1.4"
@@ -24,6 +26,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Distributions", "Documenter", "Flux", "ForwardDiff", "ReverseDiff", "Test", "StatsBase"]
+test = ["Distributions", "Documenter", "Flux", "ForwardDiff", "ReverseDiff", "Test", "StatsBase", "Zygote"]

src/AD_Zygote.jl

@@ -0,0 +1,22 @@
+import .Zygote
+
+struct ZygoteGradientLogDensity{L} <: ADGradientWrapper
+    ℓ::L
+end
+
+"""
+$(SIGNATURES)
+
+Gradient using algorithmic/automatic differentiation via Zygote.
+"""
+ADgradient(::Val{:Zygote}, ℓ::AbstractLogDensityProblem) = ZygoteGradientLogDensity(ℓ)
+
+show(io::IO, ∇ℓ::ZygoteGradientLogDensity) = print(io, "Zygote AD wrapper for ", ∇ℓ.ℓ)
+
+function logdensity(::Type{ValueGradient}, ∇ℓ::ZygoteGradientLogDensity, x::RealVector)
+    @unpack ℓ = ∇ℓ
+    y, back = Zygote.forward(_value_closure(ℓ), x)
+    isfinite(y) || return ValueGradient(y, zeros(typeof(y), length(x)))
+    grad = back(Int8(1))
+    ValueGradient(y, grad)
+end

src/LogDensityProblems.jl

@@ -250,6 +250,7 @@ function __init__()
     @require ForwardDiff="f6369f11-7733-5829-9624-2563aa707210" include("AD_ForwardDiff.jl")
     @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" include("AD_Flux.jl")
     @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" include("AD_ReverseDiff.jl")
+    @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" include("AD_Zygote.jl")
 end
 
 

test/runtests.jl

@@ -3,14 +3,18 @@ using LogDensityProblems: Value, ValueGradient
 using Test
 
 using Distributions
-import ForwardDiff, Flux, ReverseDiff
+import ForwardDiff, Flux, ReverseDiff, Zygote # for AD tests
 using Parameters: @unpack
 using DocStringExtensions: SIGNATURES
 using TransformVariables
 using Random: seed!
 
 seed!(1)
 
+####
+#### building blocks
+####
+
 """
     a ≅ b
 
@@ -123,6 +127,10 @@ end
     @test repr(∇pr) == ("ReverseDiff AD wrapper for " * repr(p))
 end
 
+####
+#### AD frameworks
+####
+
 @testset "AD via Flux" begin
     f(x) = -3*abs2(x[1])
     ℓ = TransformedLogDensity(as(Array, asℝ, 1), f)
@@ -141,6 +149,15 @@ end
     @test logdensity(ValueGradient, ∇ℓ, x) ≅ ValueGradient(f(x), -6 .* x)
 end
 
+@testset "AD via Zygote" begin
+    f(x) = -3*abs2(x[1])
+    ℓ = TransformedLogDensity(as(Array, asℝ, 1), f)
+    ∇ℓ = ADgradient(:Zygote, ℓ)
+    x = randn(1)
+    @test logdensity(Value, ℓ, x) ≅ logdensity(Value, ∇ℓ, x)
+    @test_skip logdensity(ValueGradient, ∇ℓ, x) ≅ ValueGradient(f(x), -6 .* x)
+end
+
 @testset "@iffinite" begin
     flag = [0]
     f(x) = (y = LogDensityProblems.@iffinite x; flag[1] += 1; y)