High dimensional Gaussian is slow

[xukai92]

High dimensional Gaussian in Turing is slow (> 1 hour):

using Turing

@model high_dim_gauss(D) = begin
    m ~ MvNormal(fill(0, D), 1)
end

chain = sample(high_dim_gauss(10_000), HMC(0.1, 4), 2_000, progress_style=:plain)

Corresponding AHMC + Tracker impl. is fast (~ 5.5 seconds):

using Distributions, Tracker, AdvancedHMC

D = 10_000

ℓπ(x) = logpdf(MvNormal(fill(0, D), 1), x)

function ∂ℓπ∂θ(x)
    y, back = Tracker.forward(ℓπ, x)
    return (Tracker.data(y), Tracker.data(back(1)[1]))
end

metric = UnitEuclideanMetric(D)
h = Hamiltonian(metric, ℓπ, ∂ℓπ∂θ)

samples, stats = sample(h,  StaticTrajectory(Leapfrog(0.1), 4), rand(D), 2_000; verbose=false, progress=false)

[mohamed82008]

Fixed via TuringLang/Turing.jl#1064 and TuringLang/DistributionsAD.jl#18.