mwe.jl

using StatsFuns, Distributions, Random, LazyArrays, Memoization
using Turing #master
using ReverseDiff, ForwardDiff, Zygote
using MCMCBenchmarks # I forget what we had to do to this to make it work

@model TuringModelA(nA,nB,nF,nS,a,b,f) = begin
    μ ~ Normal(1,1)
    μ_coeffs ~ MvNormal(nF,1)

    σ ~ Normal(-1,1)
    σ_coeffs ~ MvNormal(nF,0.1)

    z ~ MvNormal(nA, 1)

    q = μ .+ f*μ_coeffs .+ log1p.(exp.(σ .+ f*σ_coeffs)) .* z[a]

    b ~ arraydist(BernoulliLogit.(q))
end

lazyarray(f, x) = LazyArray(Base.broadcasted(f, x))

@model TuringModelB(nA,nB,nF,nS,a,b,f) = begin
    μ ~ Normal(1,1)
    μ_coeffs ~ MvNormal(nF,1)

    σ ~ Normal(-1,1)
    σ_coeffs ~ MvNormal(nF,0.1)

    z ~ MvNormal(nA, 1)

    q = μ .+ f*μ_coeffs .+ log1p.(exp.(σ .+ f*σ_coeffs)) .* z[a]

    b ~ arraydist(BernoulliLogit.(q))

    b ~ arraydist(lazyarray(BernoulliLogit,q))
end

TuringConfig = Turing.NUTS(1000,0.85)


function simulateData(; nA = 10, nB = 10000, nF = 5 ,kwargs...)
    a = rand(1:nA,nB)
    f = Matrix{Float64}(undef,nB,nF)
    for i in 1:nF
        f[:,i] = rand(Normal(0,1),nB)
    end
  
    μ = rand(Normal(1,1))
    μ_coeffs = rand(Normal(0,1), nF)
    σ = rand(Normal(-1,1))
    σ_coeffs = rand(Normal(0,0.1), nF)
    z = rand(Normal(0,1), nA)
  
    q = μ .+ f*μ_coeffs + log1p.(exp.(σ .+ f*σ_coeffs)) .* z[a]
    p = logistic.(q)
    
    b = rand.(Bernoulli.(p))

    return (nA=nA, nB=nB, nF=nF, nS=8, a=a, b=b, f=f)
end

Random.seed!(210548)

data = simulateData()

df = DataFrame(time=Real[],sampler=String[])

Turing.setadbackend(:zygote)

y = @timed sample(TuringModelA(data...), Turing.NUTS(500,0.85), 1000; discard_adapt=false,
progress=true)

push!(df,[y[2], "TuringZygote"])

Turing.setadbackend(:forwarddiff)

y = @timed sample(TuringModelA(data...), Turing.NUTS(500,0.85), 1000; discard_adapt=false,
progress=true)

push!(df,[y[2], "TuringForward"])

Turing.setadbackend(:reversediff)
Turing.setrdcache(true)

y = @timed sample(TuringModelA(data...), Turing.NUTS(500,0.85), 1000; discard_adapt=false,
progress=true)

push!(df,[y[2], "TuringReverse"])