Update Turing tests from 9ccd28

Author: phipsgabler

Project.toml

@@ -21,6 +21,8 @@ AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
 BinaryProvider = "b99e7846-7c00-51b0-8f62-c81ae34c0232"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+EllipticalSliceSampling = "cad2338a-1db2-11e9-3401-43bc07c9ede2"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Libtask = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -42,4 +44,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [targets]
-test = ["AdvancedHMC", "AdvancedMH", "DistributionsAD", "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"]

test/Turing/Turing.jl

@@ -9,7 +9,7 @@ module Turing
 ########################################################################
 
 using Requires, Reexport, ForwardDiff
-using Bijectors, StatsFuns, SpecialFunctions
+using DistributionsAD, Bijectors, StatsFuns, SpecialFunctions
 using Statistics, LinearAlgebra
 using Markdown, Libtask, MacroTools
 @reexport using Distributions, MCMCChains, Libtask
@@ -111,7 +111,9 @@ export  @model,                 # modelling
         VecBinomialLogit,
         OrderedLogistic,
         LogPoisson,
-        NamedDist
+        NamedDist,
+        filldist,
+        arraydist
 
 # Reexports
 using AbstractMCMC: sample, psample

test/Turing/contrib/inference/AdvancedSMCExtensions.jl

@@ -255,7 +255,7 @@ function step(model, spl::Sampler{<:IPMCMC}, VarInfos::Array{VarInfo}, is_first:
 
     # Run SMC & CSMC nodes
     for j in 1:spl.alg.n_nodes
-        VarInfos[j].num_produce = 0
+        reset_num_produce!(VarInfos[j])
         VarInfos[j] = step(model, spl.info[:samplers][j], VarInfos[j])[1]
         log_zs[j] = spl.info[:samplers][j].info[:logevidence][end]
     end

test/Turing/core/Core.jl

@@ -1,5 +1,6 @@
 module Core
 
+using DistributionsAD, Bijectors
 using MacroTools, Libtask, ForwardDiff, Random
 using Distributions, LinearAlgebra
 using ..Utilities, Reexport
@@ -10,12 +11,16 @@ using DynamicPPL: Model, runmodel!,
 using LinearAlgebra: copytri!
 using Bijectors: PDMatDistribution
 import Bijectors: link, invlink
-using DistributionsAD
 using StatsFuns: logsumexp, softmax
 @reexport using DynamicPPL
+using Requires
 
 include("container.jl")
 include("ad.jl")
+@init @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin
+    include("compat/zygote.jl")
+    export ZygoteAD
+end
 
 export  @model,
         @varname,
@@ -47,8 +52,6 @@ export  @model,
         ADBACKEND,
         setchunksize,
         verifygrad,
-        gradient_logp_forward,
-        gradient_logp_reverse,
         @varinfo,
         @logpdf,
         @sampler,

test/Turing/core/ad.jl

@@ -1,15 +1,14 @@
 ##############################
 # Global variables/constants #
 ##############################
-using Bijectors
-
 const ADBACKEND = Ref(:forward_diff)
-function setadbackend(backend_sym)
-    @assert backend_sym == :forward_diff || backend_sym == :reverse_diff
-    backend_sym == :forward_diff && CHUNKSIZE[] == 0 && setchunksize(40)
-    ADBACKEND[] = backend_sym
-
-    Bijectors.setadbackend(backend_sym)
+setadbackend(backend_sym::Symbol) = setadbackend(Val(backend_sym))
+function setadbackend(::Val{:forward_diff})
+    CHUNKSIZE[] == 0 && setchunksize(40)
+    ADBACKEND[] = :forward_diff
+end
+function setadbackend(::Val{:reverse_diff})
+    ADBACKEND[] = :reverse_diff
 end
 
 const ADSAFE = Ref(false)
@@ -39,22 +38,23 @@ ADBackend() = ADBackend(ADBACKEND[])
 ADBackend(T::Symbol) = ADBackend(Val(T))
 
 ADBackend(::Val{:forward_diff}) = ForwardDiffAD{CHUNKSIZE[]}
-ADBackend(::Val) = TrackerAD
+ADBackend(::Val{:reverse_diff}) = TrackerAD
+ADBackend(::Val) = error("The requested AD backend is not available. Make sure to load all required packages.")
 
 """
-getADtype(alg)
+    getADbackend(alg)
 
-Finds the autodifferentiation type of the algorithm `alg`.
+Find the autodifferentiation backend of the algorithm `alg`.
 """
-getADtype(spl::Sampler) = getADtype(spl.alg)
+getADbackend(spl::Sampler) = getADbackend(spl.alg)
 
 """
-gradient_logp(
-    θ::AbstractVector{<:Real},
-    vi::VarInfo,
-    model::Model,
-    sampler::AbstractSampler=SampleFromPrior(),
-)
+    gradient_logp(
+        θ::AbstractVector{<:Real},
+        vi::VarInfo,
+        model::Model,
+        sampler::AbstractSampler=SampleFromPrior(),
+    )
 
 Computes the value of the log joint of `θ` and its gradient for the model
 specified by `(vi, sampler, model)` using whichever automatic differentation
@@ -66,26 +66,23 @@ function gradient_logp(
     model::Model,
     sampler::Sampler
 )
-    ad_type = getADtype(sampler)
-    if ad_type <: ForwardDiffAD
-        return gradient_logp_forward(θ, vi, model, sampler)
-    else ad_type <: TrackerAD
-        return gradient_logp_reverse(θ, vi, model, sampler)
-    end
+    return gradient_logp(getADbackend(sampler), θ, vi, model, sampler)
 end
 
 """
-gradient_logp_forward(
+gradient_logp(
+    backend::ADBackend,
     θ::AbstractVector{<:Real},
     vi::VarInfo,
     model::Model,
-    spl::AbstractSampler=SampleFromPrior(),
+    sampler::AbstractSampler = SampleFromPrior(),
 )
 
-Computes the value of the log joint of `θ` and its gradient for the model
-specified by `(vi, spl, model)` using forwards-mode AD from ForwardDiff.jl.
+Compute the value of the log joint of `θ` and its gradient for the model
+specified by `(vi, sampler, model)` using `backend` for AD, e.g. `ForwardDiffAD{N}()` uses `ForwardDiff.jl` with chunk size `N`, `TrackerAD()` uses `Tracker.jl` and `ZygoteAD()` uses `Zygote.jl`.
 """
-function gradient_logp_forward(
+function gradient_logp(
+    ::ForwardDiffAD,
     θ::AbstractVector{<:Real},
     vi::VarInfo,
     model::Model,
@@ -110,23 +107,12 @@ function gradient_logp_forward(
 
     return l, ∂l∂θ
 end
-
-"""
-gradient_logp_reverse(
-    θ::AbstractVector{<:Real},
-    vi::VarInfo,
-    model::Model,
-    sampler::AbstractSampler=SampleFromPrior(),
-)
-
-Computes the value of the log joint of `θ` and its gradient for the model
-specified by `(vi, sampler, model)` using reverse-mode AD from Tracker.jl.
-"""
-function gradient_logp_reverse(
+function gradient_logp(
+    ::TrackerAD,
     θ::AbstractVector{<:Real},
     vi::VarInfo,
     model::Model,
-    sampler::AbstractSampler=SampleFromPrior(),
+    sampler::AbstractSampler = SampleFromPrior(),
 )
     T = typeof(getlogp(vi))
 
@@ -138,8 +124,9 @@ function gradient_logp_reverse(
 
     # Compute forward and reverse passes.
     l_tracked, ȳ = Tracker.forward(f, θ)
-    l::T, ∂l∂θ::typeof(θ) = Tracker.data(l_tracked), Tracker.data(ȳ(1)[1])
     # Remove tracking info from variables in model (because mutable state).
+    l::T, ∂l∂θ::typeof(θ) = Tracker.data(l_tracked), Tracker.data(ȳ(1)[1])
+
     return l, ∂l∂θ
 end
 
@@ -153,31 +140,7 @@ function verifygrad(grad::AbstractVector{<:Real})
     end
 end
 
-for F in (:link, :invlink)
-    @eval begin
-        function $F(
-            dist::Dirichlet, 
-            x::Tracker.TrackedArray, 
-            ::Type{Val{proj}} = Val{true}
-        ) where {proj}
-            return Tracker.track($F, dist, x, Val{proj})
-        end
-        Tracker.@grad function $F(
-            dist::Dirichlet, 
-            x::Tracker.TrackedArray, 
-            ::Type{Val{proj}}
-        ) where {proj}
-            x_data = Tracker.data(x)
-            T = eltype(x_data)
-            y = $F(dist, x_data, Val{proj})
-            return  y, Δ -> begin
-                out = (ForwardDiff.jacobian(x -> $F(dist, x, Val{proj}), x_data)::Matrix{T})' * Δ
-                return (nothing, out, nothing)
-            end
-        end
-    end
-end
-
+# These still seem necessary
 for F in (:link, :invlink)
     @eval begin
         $F(dist::PDMatDistribution, x::Tracker.TrackedArray) = Tracker.track($F, dist, x)

test/Turing/core/compat/zygote.jl

@@ -0,0 +1,29 @@
+struct ZygoteAD <: ADBackend end
+ADBackend(::Val{:zygote}) = ZygoteAD
+function setadbackend(::Val{:zygote})
+    ADBACKEND[] = :zygote
+end
+
+function gradient_logp(
+    backend::ZygoteAD,
+    θ::AbstractVector{<:Real},
+    vi::VarInfo,
+    model::Model,
+    sampler::AbstractSampler = SampleFromPrior(),
+)
+    T = typeof(getlogp(vi))
+    
+    # Specify objective function.
+    function f(θ)
+        new_vi = VarInfo(vi, sampler, θ)
+        return getlogp(runmodel!(model, new_vi, sampler))
+    end
+
+    # Compute forward and reverse passes.
+    l::T, ȳ = Zygote.pullback(f, θ)
+    ∂l∂θ::typeof(θ) = ȳ(1)[1]
+
+    return l, ∂l∂θ
+end
+
+Zygote.@nograd DynamicPPL.updategid!

test/Turing/inference/AdvancedSMC.jl

@@ -260,7 +260,7 @@ function AbstractMCMC.sample_end!(
     spl.state.average_logevidence = loge
 end
 
-function assume(spl::Sampler{<:Union{PG,SMC}}, dist::Distribution, vn::VarName, ::VarInfo)
+function DynamicPPL.assume(spl::Sampler{<:Union{PG,SMC}}, dist::Distribution, vn::VarName, ::VarInfo)
     vi = current_trace().vi
     if vn in getspace(spl)
         if ~haskey(vi, vn)
@@ -283,12 +283,13 @@ function assume(spl::Sampler{<:Union{PG,SMC}}, dist::Distribution, vn::VarName,
             r = rand(dist)
             push!(vi, vn, r, dist, Selector(:invalid))
         end
-        acclogp!(vi, logpdf_with_trans(dist, r, istrans(vi, vn)))
+        lp = logpdf_with_trans(dist, r, istrans(vi, vn))
+        acclogp!(vi, lp)
     end
     return r, 0
 end
 
-function observe(spl::Sampler{<:Union{PG,SMC}}, dist::Distribution, value, vi)
+function DynamicPPL.observe(spl::Sampler{<:Union{PG,SMC}}, dist::Distribution, value, vi)
     produce(logpdf(dist, value))
     return 0
 end