Merge pull request #114 from TuringLang/mt/update_turing_src

Update Turing src in tests

Author: mohamed82008

test/Turing/Turing.jl

@@ -11,12 +11,11 @@ module Turing
 using Requires, Reexport, ForwardDiff
 using DistributionsAD, Bijectors, StatsFuns, SpecialFunctions
 using Statistics, LinearAlgebra
-using Markdown, Libtask, MacroTools
-@reexport using Distributions, MCMCChains, Libtask
+using Libtask
+@reexport using Distributions, MCMCChains, Libtask, AbstractMCMC
 using Tracker: Tracker
 
-import Base: ~, ==, convert, hash, promote_rule, rand, getindex, setindex!
-import DynamicPPL: getspace
+import DynamicPPL: getspace, NoDist, NamedDist
 
 const PROGRESS = Ref(true)
 function turnprogress(switch::Bool)
@@ -68,6 +67,8 @@ export  @model,                 # modelling
         @varname,
         DynamicPPL,
 
+        Prior,                  # Sampling from the prior
+
         MH,                     # classic sampling
         RWMH,
         ESS,
@@ -90,7 +91,6 @@ export  @model,                 # modelling
         ADVI,
 
         sample,                 # inference
-        psample,
         setchunksize,
         resume,
         @logprob_str,
@@ -105,15 +105,10 @@ export  @model,                 # modelling
         Flat,
         FlatPos,
         BinomialLogit,
-        VecBinomialLogit,
+        BernoulliLogit,
         OrderedLogistic,
         LogPoisson,
         NamedDist,
         filldist,
         arraydist
-
-# Reexports
-using AbstractMCMC: sample, psample
-export sample, psample
-
 end

test/Turing/contrib/inference/dynamichmc.jl

@@ -46,7 +46,7 @@ mutable struct DynamicNUTSState{V<:VarInfo, D} <: AbstractSamplerState
     draws::Vector{D}
 end
 
-getspace(::DynamicNUTS{<:Any, space}) where {space} = space
+DynamicPPL.getspace(::DynamicNUTS{<:Any, space}) where {space} = space
 
 function AbstractMCMC.sample_init!(
     rng::AbstractRNG,
@@ -60,16 +60,23 @@ function AbstractMCMC.sample_init!(
         gradient_logp(x, spl.state.vi, model, spl)
     end
 
+    # Set the parameters to a starting value.
+    initialize_parameters!(spl; kwargs...)
+
     model(spl.state.vi, SampleFromUniform())
+    link!(spl.state.vi, spl)
+    l, dl = _lp(spl.state.vi[spl])
+    while !isfinite(l) || !isfinite(dl)
+        model(spl.state.vi, SampleFromUniform())
+        link!(spl.state.vi, spl)
+        l, dl = _lp(spl.state.vi[spl])
+    end
 
-    if spl.selector.tag == :default
+    if spl.selector.tag == :default && !islinked(spl.state.vi, spl)
         link!(spl.state.vi, spl)
         model(spl.state.vi, spl)
     end
 
-    # Set the parameters to a starting value.
-    initialize_parameters!(spl; kwargs...)
-
     results = mcmc_with_warmup(
         rng,
         FunctionLogDensity(
@@ -114,7 +121,7 @@ end
     model::AbstractModel,
     alg::DynamicNUTS,
     N::Integer;
-    chain_type=Chains,
+    chain_type=MCMCChains.Chains,
     resume_from=nothing,
     progress=PROGRESS[],
     kwargs...
@@ -130,19 +137,20 @@ end
     end
 end
 
-function AbstractMCMC.psample(
+function AbstractMCMC.sample(
     rng::AbstractRNG,
     model::AbstractModel,
     alg::DynamicNUTS,
+    parallel::AbstractMCMC.AbstractMCMCParallel,
     N::Integer,
     n_chains::Integer;
-    chain_type=Chains,
+    chain_type=MCMCChains.Chains,
     progress=PROGRESS[],
     kwargs...
 )
     if progress
         @warn "[$(alg_str(alg))] Progress logging in Turing is disabled since DynamicHMC provides its own progress meter"
     end
-    return AbstractMCMC.psample(rng, model, Sampler(alg, model), N, n_chains;
-                                chain_type=chain_type, progress=false, kwargs...)
+    return AbstractMCMC.sample(rng, model, Sampler(alg, model), parallel, N, n_chains;
+                               chain_type=chain_type, progress=false, kwargs...)
 end

test/Turing/contrib/inference/sghmc.jl

@@ -172,7 +172,7 @@ function step(
     spl.selector.tag != :default && link!(vi, spl)
 
     mssa = AHMC.Adaptation.ManualSSAdaptor(AHMC.Adaptation.MSSState(spl.alg.ϵ))
-    spl.info[:adaptor] = AHMC.NaiveHMCAdaptor(AHMC.UnitPreconditioner(), mssa)
+    spl.info[:adaptor] = AHMC.NaiveHMCAdaptor(AHMC.UnitMassMatrix(), mssa)
 
     spl.selector.tag != :default && invlink!(vi, spl)
     return vi, true

test/Turing/core/Core.jl

@@ -1,13 +1,12 @@
 module Core
 
 using DistributionsAD, Bijectors
-using MacroTools, Libtask, ForwardDiff, Random
+using Libtask, ForwardDiff, Random
 using Distributions, LinearAlgebra
 using ..Utilities, Reexport
 using Tracker: Tracker
 using ..Turing: Turing
-using DynamicPPL: Model,
-    AbstractSampler, Sampler, SampleFromPrior
+using DynamicPPL: Model, AbstractSampler, Sampler, SampleFromPrior
 using LinearAlgebra: copytri!
 using Bijectors: PDMatDistribution
 import Bijectors: link, invlink
@@ -17,9 +16,15 @@ using Requires
 
 include("container.jl")
 include("ad.jl")
-@init @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin
-    include("compat/zygote.jl")
-    export ZygoteAD
+function __init__()
+    @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin
+        include("compat/zygote.jl")
+        export ZygoteAD
+    end
+    @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
+        include("compat/reversediff.jl")
+        export ReverseDiffAD, getrdcache, setrdcache, emptyrdcache
+    end
 end
 
 export  @model,
@@ -36,10 +41,9 @@ export  @model,
         forkr,
         current_trace,
         getweights,
+        getweight,
         effectiveSampleSize,
-        increase_logweight,
-        inrease_logevidence,
-        resample!,
+        sweep!,
         ResampleWithESSThreshold,
         ADBackend,
         setadbackend,

test/Turing/core/ad.jl

@@ -1,14 +1,23 @@
 ##############################
 # Global variables/constants #
 ##############################
-const ADBACKEND = Ref(:forward_diff)
+const ADBACKEND = Ref(:forwarddiff)
 setadbackend(backend_sym::Symbol) = setadbackend(Val(backend_sym))
 function setadbackend(::Val{:forward_diff})
+    Base.depwarn("`Turing.setadbackend(:forward_diff)` is deprecated. Please use `Turing.setadbackend(:forwarddiff)` to use `ForwardDiff`.", :setadbackend)
+    setadbackend(Val(:forwarddiff))
+end
+function setadbackend(::Val{:forwarddiff})
     CHUNKSIZE[] == 0 && setchunksize(40)
-    ADBACKEND[] = :forward_diff
+    ADBACKEND[] = :forwarddiff
 end
+
 function setadbackend(::Val{:reverse_diff})
-    ADBACKEND[] = :reverse_diff
+    Base.depwarn("`Turing.setadbackend(:reverse_diff)` is deprecated. Please use `Turing.setadbackend(:tracker)` to use `Tracker` or `Turing.setadbackend(:reversediff)` to use `ReverseDiff`. To use `ReverseDiff`, please make sure it is loaded separately with `using ReverseDiff`.",  :setadbackend)
+    setadbackend(Val(:tracker))
+end
+function setadbackend(::Val{:tracker})
+    ADBACKEND[] = :tracker
 end
 
 const ADSAFE = Ref(false)
@@ -37,8 +46,8 @@ struct TrackerAD <: ADBackend end
 ADBackend() = ADBackend(ADBACKEND[])
 ADBackend(T::Symbol) = ADBackend(Val(T))
 
-ADBackend(::Val{:forward_diff}) = ForwardDiffAD{CHUNKSIZE[]}
-ADBackend(::Val{:reverse_diff}) = TrackerAD
+ADBackend(::Val{:forwarddiff}) = ForwardDiffAD{CHUNKSIZE[]}
+ADBackend(::Val{:tracker}) = TrackerAD
 ADBackend(::Val) = error("The requested AD backend is not available. Make sure to load all required packages.")
 
 """

test/Turing/core/compat/reversediff.jl

@@ -0,0 +1,93 @@
+using .ReverseDiff: compile, GradientTape
+using .ReverseDiff.DiffResults: GradientResult
+
+struct ReverseDiffAD{cache} <: ADBackend end
+const RDCache = Ref(false)
+setrdcache(b::Bool) = setrdcache(Val(b))
+setrdcache(::Val{false}) = RDCache[] = false
+setrdcache(::Val) = throw("Memoization.jl is not loaded. Please load it before setting the cache to true.")
+function emptyrdcache end
+
+getrdcache() = RDCache[]
+ADBackend(::Val{:reversediff}) = ReverseDiffAD{getrdcache()}
+function setadbackend(::Val{:reversediff})
+    ADBACKEND[] = :reversediff
+end
+
+function gradient_logp(
+    backend::ReverseDiffAD{false},
+    θ::AbstractVector{<:Real},
+    vi::VarInfo,
+    model::Model,
+    sampler::AbstractSampler = SampleFromPrior(),
+)
+    T = typeof(getlogp(vi))
+    
+    # Specify objective function.
+    function f(θ)
+        new_vi = VarInfo(vi, sampler, θ)
+        model(new_vi, sampler)
+        return getlogp(new_vi)
+    end
+    tp, result = taperesult(f, θ)
+    ReverseDiff.gradient!(result, tp, θ)
+    l = DiffResults.value(result)
+    ∂l∂θ::typeof(θ) = DiffResults.gradient(result)
+
+    return l, ∂l∂θ
+end
+
+tape(f, x) = GradientTape(f, x)
+function taperesult(f, x)
+    return tape(f, x), GradientResult(x)
+end
+
+@require Memoization = "6fafb56a-5788-4b4e-91ca-c0cea6611c73" @eval begin
+    setrdcache(::Val{true}) = RDCache[] = true
+    function emptyrdcache()
+        for k in keys(Memoization.caches)
+            if k[1] === typeof(memoized_taperesult)
+                pop!(Memoization.caches, k)
+            end
+        end
+    end
+    function gradient_logp(
+        backend::ReverseDiffAD{true},
+        θ::AbstractVector{<:Real},
+        vi::VarInfo,
+        model::Model,
+        sampler::AbstractSampler = SampleFromPrior(),
+    )
+        T = typeof(getlogp(vi))
+        
+        # Specify objective function.
+        function f(θ)
+            new_vi = VarInfo(vi, sampler, θ)
+            model(new_vi, sampler)
+            return getlogp(new_vi)
+        end
+        ctp, result = memoized_taperesult(f, θ)
+        ReverseDiff.gradient!(result, ctp, θ)
+        l = DiffResults.value(result)
+        ∂l∂θ = DiffResults.gradient(result)
+
+        return l, ∂l∂θ
+    end
+
+    # This makes sure we generate a single tape per Turing model and sampler
+    struct RDTapeKey{F, Tx}
+        f::F
+        x::Tx
+    end
+    function Memoization._get!(f::Union{Function, Type}, d::IdDict, keys::Tuple{Tuple{RDTapeKey}, Any})
+        key = keys[1][1]
+        return Memoization._get!(f, d, (typeof(key.f), typeof(key.x), size(key.x)))
+    end
+    memoized_taperesult(f, x) = memoized_taperesult(RDTapeKey(f, x))
+    Memoization.@memoize function memoized_taperesult(k::RDTapeKey)
+        return compiledtape(k.f, k.x), GradientResult(k.x)
+    end
+    memoized_tape(f, x) = memoized_tape(RDTapeKey(f, x))
+    Memoization.@memoize memoized_tape(k::RDTapeKey) = compiledtape(k.f, k.x)
+    compiledtape(f, x) = compile(GradientTape(f, x))
+end