JuliaDiff / ForwardDiff.jl

Forward Mode Automatic Differentiation for Julia
Other
888 stars 141 forks source link

Implement gammalogccdf for ForwardDiff #666

Open mashu opened 12 months ago

mashu commented 12 months ago

I am modeling in Turing.jl some process in which I use Gamma prior to parametrize Truncated Poisson and I use NUTs sampler.

However, my model has nested dependence and I cannot fully vectorize the loop, so Zygote would be slow. The default ForwardDiff throws at me an error

ERROR: MethodError: no method matching _gammalogccdf(::ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 8}, ::ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 8}, ::ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 8})
Stacktrace:
  [1] gammalogccdf(k::Int64, θ::Int64, x::ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 8})
    @ StatsFuns ~/.julia/packages/StatsFuns/atA5S/src/distrs/gamma.jl:71
  [2] poislogcdf(λ::ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 8}, x::Int64)
    @ StatsFuns ~/.julia/packages/StatsFuns/atA5S/src/distrs/pois.jl:30
  [3] logcdf(d::Poisson{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 8}}, x::Int64)
    @ Distributions ~/.julia/packages/Distributions/Ufrz2/src/univariates.jl:638
  [4] _logcdf_noninclusive(d::Poisson{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 8}}, x::Int64)
    @ Distributions ~/.julia/packages/Distributions/Ufrz2/src/censored.jl:441
  [5] truncated(d::Poisson{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 8}}, l::Int64, u::Int64)
    @ Distributions ~/.julia/packages/Distributions/Ufrz2/src/truncate.jl:67
  [6] Truncated(d::Poisson{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 8}}, l::Int64, u::Int64)
    @ Distributions ./deprecated.jl:105
  [7] macro expansion

To my understanding this error suggests that the model is encountering issues with automatic differentiation, specifically when evaluating the complementary cumulative distribution function (CCDF) of the Gamma distribution. Seems like the error might be related to the Truncated Poisson distribution, as the truncation operation might involve the CCDF of the underlying distribution.

Can that be solved in ForwardDiff AD ?

devmotion commented 12 months ago

I think this issue should be filed with StatsFuns. The missing part is support of ForwardDiff for StatsFuns.gammalogccdf and/or StatsFuns.poislogcdf. These functions do not support ForwardDiff.Dual yet. The problem is similar to https://github.com/JuliaStats/StatsFuns.jl/issues/152.