Update ChainRules definitions and add differential for PoissonBinomial pdf

[devmotion]

The PR updates the ChainRules definitions and tests, according to JuliaStats/StatsFuns.jl#106 (which ideally would allow us to remove the definitions in DistributionsAD).

Additionally, I derived and implemented the differential for the pdf of the PoissonBinomial distribution. It uses a similar dynamic programming style as the implementation of the pdf in Distributions. Unfortunately, the recently added lazy initialization of PoissonBinomial (JuliaStats/Distributions.jl#1285) breaks Zygote support. I didn't manage to fix it with a custom differential for getproperty (or the Zygote analogue literal_getproperty). E.g., the following caused errors

    # this rrule is necessary since the primal mutates
    function ChainRulesCore.rrule(
        ::typeof(getproperty),
        d::T,
        x::Symbol,
    ) where {T<:PoissonBinomial}
        y, A = if getfield(d, :pmf) === nothing && x === :pmf
            getproperty(d, x), poissonbinomial_partialderivatives(d.p)
        else
            getproperty(d, x), nothing
        end
        
        function getproperty_PoissonBinomial_pullback(Δy)
            ∂d = if x === :pmf && A !== nothing
                Composite{T}(; pmf=A * Δy)
            else
                DoesNotExist()
            end
            return NO_FIELDS, ∂d, DoesNotExist()
        end

        return y, getproperty_PoissonBinomial_pullback
    end

similar to those reported in FluxML/Zygote.jl#566. I think this issue needs more time and therefore I marked the Zygote tests of PoissonBinomial as broken.

[torfjelde]

LGTM!
I haven't checked all the derivatives in detail, but that's what the tests are for 🙃

I have a couple of questions, but they're mostly just questions; don't think they'll require any change, so I just approve now 👍

[torfjelde]

Where is logtwo from? I see it's from StatsFuns, but I can't find the definition in StatsFuns 😕

I ask because I worry it might lead to undeseriable type-promotion, e.g. Float32 to Float64.

[devmotion]

It's defined in LogExpFunctions and just reexported from StatsFuns (the PR was finally merged some days ago 🎉). It is defined as an Irrational and hence should avoid type promotions if possible, e.g:

julia> logtwo - 3.4
-2.7068528194400545

julia> logtwo - 3.4f0
-2.706853f0

[torfjelde]

Good, good 👍 Yeah in this case it should be fine, but there are several cases in DiffRules.jl (and I'm pretty certain the same ones exist in ChainRules.jl?) where the operations are written in a way such that Irrational is first converted into Float64, and thus the dual is promoted to Float64, despite the primal being Float32 (JuliaDiff/DiffRules.jl#55).

And the annoying bit is that some of these constants are defined in StatsFuns.jl, some in LogExpFunctions.jl, etc., so it's difficult to re-use them in DiffRules.jl 😕
Though maybe I should actually just add all those constants to that PR instead of doing oftype everywhere, hmm.

[devmotion]

It's only a problem if you perform operations with the constants (such as sqrt) before they are promoted to the correct type, which is not the case here.

[devmotion]

I haven't checked all the derivatives in detail, but that's what the tests are for

On a side note, the testing utilities in ChainRulesTestUtils are really great (and it's also documented now), it makes it very easy to check the differentials and to ensure that they are correct. And it even warns you if e.g. you use a thunk but there is only one derivative.