Make sure all distributions support AD where possible

[mohamed82008]

The following link shows the list of distributions that StatsFuns uses ccall for: https://github.com/JuliaStats/StatsFuns.jl/blob/7473d0d8f7c5b2f4a53f0625ada5f2323c84ffe4/src/rmath.jl#L99. We should make AD workarounds for all of these.

Supersedes #673

[xukai92]

There are some distributions in this list which already have callback functions, e.g. beta (https://github.com/JuliaStats/StatsFuns.jl/blob/7473d0d8f7c5b2f4a53f0625ada5f2323c84ffe4/src/distrs/beta.jl). I actually implemented those in StatsFun for the AD purpose in Turing so I'm happy to work on this issue on the StatsFuns side (they gave quite good review on numerical stability).

I think you also implemented some on the Turing side, e.g. https://github.com/TuringLang/Turing.jl/pull/664/files. Maybe we should port them to StatsFuns (and DiffRules) as well. How do you think?

[mohamed82008]

@xukai92 Yes, StatsFuns or DiffRules with StatsFuns in a require block are the more stable solutions here (no type piracy).

[yebai]

Here is the ful list of distributions that requires special treatments for AD:

• beta beta α β Adding pure Julia pdf and logpdf as callbacks. JuliaStats/StatsFuns.jl#21
• binom binom n p Enable backprop through binomlogpdf #601
• chisq chisq k Adding pure Julia pdf and logpdf as callbacks. JuliaStats/StatsFuns.jl#21
• fdist d1 d2 Adding pure Julia pdf and logpdf as callbacks. JuliaStats/StatsFuns.jl#21
• gamma gamma α β Adding pure Julia pdf and logpdf as callbacks. JuliaStats/StatsFuns.jl#21
• hyper hyper ms mf n
• norm norm μ σ Adding pure Julia pdf and logpdf as callbacks. JuliaStats/StatsFuns.jl#21
• nbinom nbinom r p AD workaround for nbinomlogpdf #664
• pois pois λ 699b9db
• t tdist k Adding pure Julia pdf and logpdf as callbacks. JuliaStats/StatsFuns.jl#21
• nbeta nbeta α β λ
• nchisq nchisq k λ
• nfdist k1 k2 λ
• ntdist k λ

Related: JuliaStats/StatsFuns.jl#60 #615

[xukai92]

We also need to test both forward and reverse mode (different AD implementation). And vector operations (different in bijection mapping) for all.

[yebai]

Currently, the following distributions need a generic Julia implementation for pdf and logpdf,

• hyper hyper ms mf n (discrete)
• nbeta nbeta α β λ (obsolete, not exported by Distributions.jl)
• nchisq nchisq k λ (obsolete, not exported by Distributions.jl)
• nfdist k1 k2 λ (obsolete, not exported by Distributions.jl)
• ntdist k λ (obsolete, not exported by Distributions.jl)

This can be done in a similar way to JuliaStats/StatsFuns.jl#21, or #664. We should also consider moving the hotfixes for pois, nbinom and binom into StatsFuns.jl/src/distrs/.

[yebai]

Here's the full list of distributions from Distributions.jl, but without FluxAD support:

• MvNormal MvNormal and Tracker #717
• UniveriateGMM AD fails for MixtureModel #722

[mohamed82008]

This seems to be fixed now. I will close the issue.