Add basic AD test

Author: devmotion

Project.toml

@@ -45,6 +45,7 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["AdvancedHMC", "AdvancedMH", "DistributionsAD", "DocStringExtensions", "EllipticalSliceSampling", "ForwardDiff", "Libtask", "LinearAlgebra", "LogDensityProblems", "Logging", "MCMCChains", "Markdown", "PDMats", "ProgressLogging", "Random", "Reexport", "Requires", "SpecialFunctions", "Statistics", "StatsBase", "StatsFuns", "Test", "Tracker", "UUIDs"]
+test = ["AdvancedHMC", "AdvancedMH", "DistributionsAD", "DocStringExtensions", "EllipticalSliceSampling", "ForwardDiff", "Libtask", "LinearAlgebra", "LogDensityProblems", "Logging", "MCMCChains", "Markdown", "PDMats", "ProgressLogging", "Random", "Reexport", "Requires", "SpecialFunctions", "Statistics", "StatsBase", "StatsFuns", "Test", "Tracker", "UUIDs", "Zygote"]

test/compat/ad.jl

@@ -0,0 +1,67 @@
+using DynamicPPL
+using Distributions
+
+using ForwardDiff
+using Zygote
+using Tracker
+
+@testset "logp" begin
+    @model function admodel()
+        s ~ InverseGamma(2, 3)
+        m ~ Normal(0, sqrt(s))
+        1.5 ~ Normal(m, sqrt(s))
+        2.0 ~ Normal(m, sqrt(s))
+        return s, m
+    end
+
+    model = admodel()
+    vi = VarInfo(model)
+    model(vi, SampleFromPrior())
+    x = [vi[@varname(s)], vi[@varname(m)]]
+
+    dist_s = InverseGamma(2,3)
+
+    # Hand-written log probabilities for vector `x = [s, m]`.
+    function logp_manual(x)
+        s = x[1]
+        m = x[2]
+        dist = Normal(m, sqrt(s))
+
+        return logpdf(InverseGamma(2, 3), s) + logpdf(Normal(0, sqrt(s)), m) +
+            logpdf(dist, 1.5) + logpdf(dist, 2.0)
+    end
+
+    # Log probabilities for vector `x = [s, m]` using the model.
+    function logp_model(x)
+        new_vi = VarInfo(vi, SampleFromPrior(), x)
+        model(new_vi, SampleFromPrior())
+        return getlogp(new_vi)
+    end
+
+    # Check that both functions return the same values.
+    lp = logp_manual(x)
+    @test logp_model(x) ≈ lp
+
+    # Gradients based on the manual implementation.
+    grad = ForwardDiff.gradient(logp_manual, x)
+
+    y, back = Tracker.forward(logp_manual, x)
+    @test Tracker.data(y) ≈ lp
+    @test Tracker.data(back(1)[1]) ≈ grad
+
+    y, back = Zygote.pullback(logp_manual, x)
+    @test y ≈ lp
+    @test back(1)[1] ≈ grad
+
+    # Gradients based on the model.
+    @test ForwardDiff.gradient(logp_model, x) ≈ grad
+
+    y, back = Tracker.forward(logp_model, x)
+    @test Tracker.data(y) ≈ lp
+    @test Tracker.data(back(1)[1]) ≈ grad
+
+    y, back = Zygote.gradient(logp_model, x)
+    @test y ≈ lp
+    @test back(1) ≈ grad
+end
+

test/runtests.jl

@@ -8,6 +8,7 @@ turnprogress(false)
 @testset "DynamicPPL.jl" begin
     include("utils.jl")
     include("compiler.jl")
+    include("compat/ad.jl")
     include("varinfo.jl")
     include("sampler.jl")
     include("prob_macro.jl")