Help -- can't get Turing to work on a simple model

Specific Domains Probabilistic Programming
,
2

Hello @anprasad!

I don’t think this is user error at all! It is a bug :scream: Specifically, it is a bug with automatic differentiation. If you use a gradientless sampler (e.g. MH()) it runs fine.

Here’s a simplified example, which illustrates the bug in that the gradients contain NaN’s, and unfortunately none of our supported AD backends are happy with this:

using DynamicPPL: @model, LogDensityFunction
using Distributions
using LogDensityProblems: logdensity_and_gradient
using LogDensityProblemsAD: ADgradient

@model function model1()
    σ ~ InverseGamma(2, 3)
    V ~ truncated(Normal(0, σ), 0, Inf)
end

import ForwardDiff
ℓ = ADgradient(:ForwardDiff, LogDensityFunction(model1()))
logdensity_and_gradient(ℓ, [1.0, 2.0])
# (-3.0285667753085077, [NaN, NaN])

import Mooncake
ℓ = ADgradient(:Mooncake, LogDensityFunction(model1()))
logdensity_and_gradient(ℓ, [1.0, 2.0, 3.0])
# (-3.0285667753085077, [NaN, -2.0])

import ReverseDiff
ℓ = ADgradient(:ReverseDiff, LogDensityFunction(model1()))
logdensity_and_gradient(ℓ, [1.0, 2.0, 3.0])
# (-3.0285667753085077, [NaN, -2.0])

import Zygote
ℓ = ADgradient(:Zygote, LogDensityFunction(model1()))
logdensity_and_gradient(ℓ, [1.0, 2.0, 3.0])
# (-3.0285667753085077, [NaN, -2.0])

However: for some reason which is yet unclear, it works when I replace

truncated(Normal(a, b), 0, Inf)

with

truncated(Normal(a, b); lower=0)

in your model. This has the same meaning, but somehow the AD packages are happy with it. And as it happens, this change only needs to be made in the definition of V:

using Turing

@model function model1(samples)
    V0 = 100
    σ ~ truncated(Normal(0.01 * V0, 0.15 * V0), 0, Inf)
    C0 ~ truncated(Normal(100, 60), 0, Inf)
    V ~ filldist(truncated(Normal(V0, σ); lower=0), length(samples))
    for i in eachindex(samples)
        samples[i] ~ Poisson(C0 * V[i])
    end
end

samples = [1, 2, 3]

sample(model1(samples), NUTS(), 1000)

So perhaps that could be a good enough workaround while we try to get the AD bugs fixed? :smile:

If it makes a difference, I’m running this on macOS Julia 1.11.1, with Turing v0.35.1.

Thank you for posting this!

2 Likes