Can Turing differentiate a truncated Poisson

General Usage
, ,
1

does Turing’s AD backend struggle with truncated Poisson distributions?

Turing.@model function bps_hurdle_model(bonus::Vector{Int}, bps::Vector{Int}, 
    μ_αbps = 0, σ_αbps = 1, μ_βbps = 0, σ_βbps = 1)

    α_pzero ~ Normal(-1, 1); β_pzero ~ Normal(1/3, 1)
    α_bps ~ Normal(μ_αbps, σ_αbps); β_bps ~ Normal(μ_βbps, σ_βbps)
    
    for n ∈ 1:length(bonus)
        p_zero = logistic(α_pzero + β_pzero * bps[n])
        λ = exp(α_bps + β_bps * bps[n])
        
        if bonus[n] == 0
            Turing.@addlogprob! log(p_zero)
        else
            Turing.@addlogprob! log(1 - p_zero)
            bonus[n] ~ truncated(Poisson(λ), 1, 3)
        end
    end
end

this model results in the following error:

julia> bps_hurdle_model(Int.(bonus_data.bonus), Int.(bonus_data.bps)) |>
           model -> sample(model, NUTS(), MCMCThreads(), n_draws, n_chains)
ERROR: TaskFailedException

    nested task error: TaskFailedException
    
        nested task error: MethodError: no method matching _gammalogccdf(::ForwardDiff.Dual{…}, ::ForwardDiff.Dual{…}, ::ForwardDiff.Dual{…})
        The function `_gammalogccdf` exists, but no method is defined for this combination of argument types.
        
        Closest candidates are:
          _gammalogccdf(::Float64, ::Float64, ::Float64)
           @ StatsFuns ~/.julia/packages/StatsFuns/NSRpu/src/distrs/gamma.jl:80
          _gammalogccdf(::T, ::T, ::T) where T<:Union{Float16, Float32}
           @ StatsFuns ~/.julia/packages/StatsFuns/NSRpu/src/distrs/gamma.jl:97

…despite there being no gamma distributions explictly defined in the model.

switching to non-gradient based samplers for the bonus parameters allows the model to run, but not very successfully.

I have an alternative parameterisation that seems to run using NUTS(), but I do not understand the original issue. If Stan can sample from an equivalent model where I use a truncated Poisson distribution, what is causing this error?

2

The _gammalogccdf is used in computing the logcdf of the Poisson distribution: julia> @less StatsFuns.poislogcdf(1.2, 3).
Unfortunately, it is not generic enough to support the dual number type used by ForwardDiff. You can try using a different AD backend which works does not pass around special number types, but works by transforming the source code directly. Seems like Turing supports Zygote and Mooncake which are both able to take the derivative of the truncated Poisson distribution.

If Stan can sample from an equivalent model where I use a truncated Poisson distribution, what is causing this error?

Stan is a closed, controlled environment were every defined function plays nicely with its AD. Julia is more open in that packages can be combined freely and pull in arbitrary functions. Depending on how those are written, i.e., pure Julia or calling C in the end, they might or might not be supported by a particular AD backend. Luckily, it’s easy to switch the AD backend in Julia and most code can be made to work.

1 Like