Fix test errors and update to ChainRules 1 and Distributions 0.25.15 (#…

…198)

* Fix warning

* Mark BetaPrime as broken

* Test Julia 1.3 only with Ubuntu

* Bump version

* Not broken anymore?!

* Fix numerical issues of BetaPrime for α=1

* Update to CR1

* Remove NNlib

* Fix error

* Add rrule for to_simplex

* Prepare for addition of ChainRules definitions to StatsFuns

* Prepare for use of CR in Distributions

* Add Zygote workaround

* Handle Zygote's `nothing` in tests

* Remove ChainRules definitions that are moved to StatsFuns

* Remove ChainRules definition that was moved to Distributions

* Fix deprecation

* Fix Tracker gradient of `poissonbinomial_pdf`

* Add Zygote workarounds

* Update list of broken distributions

* Remove type piracy (default definition is sufficient)

* Use `eachcol`

.github/workflows/AD.yml

@@ -17,14 +17,30 @@ jobs:
           - '1'
         os:
           - ubuntu-latest
-          - macOS-latest
         arch:
           - x64
         AD:
           - ForwardDiff
           - Tracker
           - ReverseDiff
           - Zygote
+        include:
+          - version: '1'
+            os: macOS-latest
+            arch: x64
+            AD: ForwardDiff
+          - version: '1'
+            os: macOS-latest
+            arch: x64
+            AD: Tracker
+          - version: '1'
+            os: macOS-latest
+            arch: x64
+            AD: ReverseDiff
+          - version: '1'
+            os: macOS-latest
+            arch: x64
+            AD: Zygote
     steps:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@v1

.github/workflows/Others.yml

@@ -16,9 +16,12 @@ jobs:
           - '1'
         os:
           - ubuntu-latest
-          - macOS-latest
         arch:
           - x64
+        include:
+          - version: '1'
+            os: macOS-latest
+            arch: x64
     steps:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@v1

Project.toml

@@ -1,6 +1,6 @@
 name = "DistributionsAD"
 uuid = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
-version = "0.6.29"
+version = "0.6.30"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -23,18 +23,18 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
 Adapt = "2, 3"
-ChainRules = "0.7, 0.8"
-ChainRulesCore = "0.9.44, 0.10"
+ChainRules = "1"
+ChainRulesCore = "1"
 Compat = "3.6"
 DiffRules = "0.1, 1.0"
-Distributions = "0.23.3, 0.24, 0.25"
+Distributions = "0.25.15"
 FillArrays = "0.8, 0.9, 0.10, 0.11"
 NaNMath = "0.3"
 PDMats = "0.9, 0.10, 0.11"
 Requires = "1"
 SpecialFunctions = "0.8, 0.9, 0.10, 1"
 StaticArrays = "0.12, 1.0"
 StatsBase = "0.32, 0.33"
-StatsFuns = "0.8, 0.9"
+StatsFuns = "0.9.10"
 ZygoteRules = "0.2"
 julia = "1.3"

src/DistributionsAD.jl

@@ -45,9 +45,6 @@ export TuringScalMvNormal,
        arraydist,
        filldist
 
-# check if Distributions >= 0.24 by checking if a generic implementation of `pdf` is defined
-const DISTRIBUTIONS_HAS_GENERIC_UNIVARIATE_PDF = hasmethod(pdf, Tuple{UnivariateDistribution,Real}) 
-
 include("common.jl")
 include("arraydist.jl")
 include("filldist.jl")

src/arraydist.jl

@@ -6,13 +6,6 @@ function arraydist(dists::AbstractVector{<:UnivariateDistribution})
     return Product(dists)
 end
 
-function Distributions.logpdf(dist::VectorOfUnivariate, x::AbstractMatrix{<:Real})
-    size(x, 1) == length(dist) ||
-        throw(DimensionMismatch("Inconsistent array dimensions."))
-    # `eachcol` breaks Zygote, so we use `view` directly
-    return map(i -> sum(map(logpdf, dist.v, view(x, :, i))), axes(x, 2))
-end
-
 struct MatrixOfUnivariate{
     S <: ValueSupport,
     Tdist <: UnivariateDistribution{S},
@@ -56,8 +49,7 @@ function arraydist(dists::AbstractVector{<:MultivariateDistribution})
 end
 
 function Distributions._logpdf(dist::VectorOfMultivariate, x::AbstractMatrix{<:Real})
-    # `eachcol` breaks Zygote, so we use `view` directly
-    return sum(i -> logpdf(dist.dists[i], view(x, :, i)), axes(x, 2))
+    return sum(((di, xi),) -> logpdf(di, xi), zip(dist.dists, eachcol(x)))
 end
 function Distributions.logpdf(dist::VectorOfMultivariate, x::AbstractArray{<:AbstractMatrix{<:Real}})
     return map(x -> logpdf(dist, x), x)

src/chainrules.jl

@@ -6,136 +6,6 @@
         insupport = a <= x <= b,
         diff = b - a,
         c = insupport ? inv(diff) : inv(one(diff)),
-        z = insupport ? zero(x) : oftype(x, NaN),
     ),
-    (c, -c, z),
+    (c, -c, ZeroTangent()),
 )
-
-# StatsFuns: https://github.com/JuliaStats/StatsFuns.jl/pull/106
-
-## Beta ##
-
-@scalar_rule(
-    betalogpdf(α::Real, β::Real, x::Number),
-    @setup(z = digamma(α + β)),
-    (
-        log(x) + z - digamma(α),
-        log1p(-x) + z - digamma(β),
-        (α - 1) / x + (1 - β) / (1 - x),
-    ),
-)
-
-## Gamma ##
-
-@scalar_rule(
-    gammalogpdf(k::Real, θ::Real, x::Number),
-    @setup(
-        invθ = inv(θ),
-        xoθ = invθ * x,
-        z = xoθ - k,
-    ),
-    (
-        log(xoθ) - digamma(k),
-        invθ * z,
-        - (1 + z) / x,
-    ),
-)
-
-## Chisq ##
-
-@scalar_rule(
-    chisqlogpdf(k::Real, x::Number),
-    @setup(hk = k / 2),
-    (
-        (log(x) - logtwo - digamma(hk)) / 2,
-        (hk - 1) / x - one(hk) / 2,
-    ),
-)
-
-## FDist ##
-
-@scalar_rule(
-    fdistlogpdf(ν1::Real, ν2::Real, x::Number),
-    @setup(
-        xν1 = x * ν1,
-        temp1 = xν1 + ν2,
-        a = (x - 1) / temp1,
-        νsum = ν1 + ν2,
-        di = digamma(νsum / 2),
-    ),
-    (
-        (-log1p(ν2 / xν1) - ν2 * a + di - digamma(ν1 / 2)) / 2,
-        (-log1p(xν1 / ν2) + ν1 * a + di - digamma(ν2 / 2)) / 2,
-        ((ν1 - 2) / x - ν1 * νsum / temp1) / 2,
-    ),
-)
-
-## TDist ##
-
-@scalar_rule(
-    tdistlogpdf(ν::Real, x::Number),
-    @setup(
-        νp1 = ν + 1,
-        xsq = x^2,
-        invν = inv(ν),
-        a = xsq * invν,
-        b = νp1 / (ν + xsq),
-    ),
-    (
-        (digamma(νp1 / 2) - digamma(ν / 2) + a * b - log1p(a) - invν) / 2,
-        - x * b,
-    ),
-)
-
-## Binomial ##
-
-@scalar_rule(
-    binomlogpdf(n::Real, p::Real, k::Real),
-    @setup(z = digamma(n - k + 1)),
-    (
-        digamma(n + 2) - z + log1p(-p) - 1 / (1 + n),
-        (k / p - n) / (1 - p),
-        z - digamma(k + 1) + logit(p),
-    ),
-)
-
-## Poisson ##
-
-@scalar_rule(
-    poislogpdf(λ::Number, x::Number),
-    ((iszero(x) && iszero(λ) ? zero(x / λ) : x / λ) - 1, log(λ) - digamma(x + 1)),
-)
-
-## PoissonBinomial
-
-function ChainRulesCore.rrule(
-    ::typeof(Distributions.poissonbinomial_pdf_fft), p::AbstractVector{<:Real}
-)
-    y = Distributions.poissonbinomial_pdf_fft(p)
-    A = poissonbinomial_partialderivatives(p)
-    function poissonbinomial_pdf_fft_pullback(Δy)
-        p̄ = InplaceableThunk(
-            @thunk(A * Δy),
-            Δ -> LinearAlgebra.mul!(Δ, A, Δy, true, true),
-        )
-        return (NO_FIELDS, p̄)
-    end
-    return y, poissonbinomial_pdf_fft_pullback
-end
-
-if isdefined(Distributions, :poissonbinomial_pdf)
-    function ChainRulesCore.rrule(
-        ::typeof(Distributions.poissonbinomial_pdf), p::AbstractVector{<:Real}
-    )
-        y = Distributions.poissonbinomial_pdf(p)
-        A = poissonbinomial_partialderivatives(p)
-        function poissonbinomial_pdf_pullback(Δy)
-            p̄ = InplaceableThunk(
-                @thunk(A * Δy),
-                Δ -> LinearAlgebra.mul!(Δ, A, Δy, true, true),
-            )
-            return (NO_FIELDS, p̄)
-        end
-        return y, poissonbinomial_pdf_pullback
-    end
-end

src/common.jl

@@ -46,45 +46,3 @@ parameterless_type(x) = parameterless_type(typeof(x))
 parameterless_type(x::Type) = __parameterless_type(x)
 
 @non_differentiable adapt_randn(::Any...)
-
-# PoissonBinomial
-
-# compute matrix of partial derivatives [∂P(X=j-1)/∂pᵢ]_{i=1,…,n; j=1,…,n+1}
-#
-# This uses the same dynamic programming "trick" as for the computation of the primals
-# in Distributions
-#
-# Reference (for the primal):
-#
-#      Marlin A. Thomas & Audrey E. Taub (1982)
-#      Calculating binomial probabilities when the trial probabilities are unequal,
-#      Journal of Statistical Computation and Simulation, 14:2, 125-131, DOI: 10.1080/00949658208810534
-function poissonbinomial_partialderivatives(p)
-    n = length(p)
-    A = zeros(eltype(p), n, n + 1)
-    @inbounds for j in 1:n
-        A[j, end] = 1
-    end
-    @inbounds for (i, pi) in enumerate(p)
-        qi = 1 - pi
-        for k in (n - i + 1):n
-            kp1 = k + 1
-            for j in 1:(i - 1)
-                A[j, k] = pi * A[j, k] + qi * A[j, kp1]
-            end
-            for j in (i+1):n
-                A[j, k] = pi * A[j, k] + qi * A[j, kp1]
-            end
-        end
-        for j in 1:(i-1)
-            A[j, end] *= pi
-        end
-        for j in (i+1):n
-            A[j, end] *= pi
-        end
-    end
-    @inbounds for j in 1:n, i in 1:n
-        A[i, j] -= A[i, j+1]
-    end
-    return A
-end

src/forwarddiff.jl

@@ -48,19 +48,3 @@ function nbinomlogpdf(r::ForwardDiff.Dual{T}, p::Real, k::Int) where {T}
     Δ_r = ForwardDiff.partials(r) * _nbinomlogpdf_grad_1(val_r, p, k)
     return FD(nbinomlogpdf(val_r, p, k),  Δ_r)
 end
-
-## ForwardDiff broadcasting support ##
-# If we use Distributions >= 0.24, then `DISTRIBUTIONS_HAS_GENERIC_UNIVARIATE_PDF` is `true`.
-# In Distributions 0.24 `logpdf` is defined for inputs of type `Real` which are then
-# converted to the support of the distributions (such as integers) in their concrete implementations.
-# Thus it is no needed to have a special function for dual numbers that performs the conversion
-# (and actually this method leads to method ambiguity errors since even discrete distributions now
-# define logpdf(::MyDistribution, ::Real), see, e.g.,
-# JuliaStats/Distributions.jl@ae2d6c5/src/univariate/discrete/binomial.jl#L119).
-if !DISTRIBUTIONS_HAS_GENERIC_UNIVARIATE_PDF
-    @eval begin
-        function Distributions.logpdf(d::DiscreteUnivariateDistribution, k::ForwardDiff.Dual)
-            return logpdf(d, convert(Integer, ForwardDiff.value(k)))
-        end
-    end
-end

src/matrixvariate.jl

@@ -214,24 +214,3 @@ function Distributions._rand!(rng::AbstractRNG, d::TuringInverseWishart, A::Abst
     X = Distributions._rand!(rng, TuringWishart(d.df, inv(cholesky(d.S))), A)
     A .= inv(cholesky!(X))
 end
-
-# Only needed in Distributions < 0.24
-if !DISTRIBUTIONS_HAS_GENERIC_UNIVARIATE_PDF
-    for T in (:MatrixBeta, :MatrixNormal, :Wishart, :InverseWishart,
-              :TuringWishart, :TuringInverseWishart,
-              :VectorOfMultivariate, :MatrixOfUnivariate)
-        @eval begin
-            Distributions.loglikelihood(d::$T, X::AbstractMatrix{<:Real}) = logpdf(d, X)
-            function Distributions.loglikelihood(d::$T, X::AbstractArray{<:Real,3})
-                (size(X, 1), size(X, 2)) == size(d) || throw(DimensionMismatch("Inconsistent array dimensions."))
-                return sum(i -> _logpdf(d, view(X, :, :, i)), axes(X, 3))
-            end
-            function Distributions.loglikelihood(
-                d::$T,
-                X::AbstractArray{<:AbstractMatrix{<:Real}},
-            )
-                return sum(x -> logpdf(d, x), X)
-            end
-        end
-    end
-end

src/reversediff.jl

@@ -18,7 +18,7 @@ using ..DistributionsAD: DistributionsAD
 
 
 import SpecialFunctions, NaNMath
-import ..DistributionsAD: turing_chol, symm_turing_chol, _mv_categorical_logpdf, adapt_randn,
+import ..DistributionsAD: turing_chol, symm_turing_chol, adapt_randn,
     simplex_logpdf
 import Base.Broadcast: materialize
 import StatsFuns: logsumexp

src/tracker.jl

@@ -261,23 +261,17 @@ end
 PoissonBinomial(p::TrackedArray{<:Real}; check_args=true) =
     TuringPoissonBinomial(p; check_args = check_args)
 
-poissonbinomial_pdf_fft(x::TrackedArray) = track(poissonbinomial_pdf_fft, x)
-@grad function poissonbinomial_pdf_fft(x::TrackedArray)
-    x_data = data(x)
-    value = poissonbinomial_pdf_fft(x_data)
-    A = poissonbinomial_partialderivatives(x_data)
-    poissonbinomial_pdf_fft_pullback(Δ) = (A * Δ,)
-    return value, poissonbinomial_pdf_fft_pullback
-end
-
-if isdefined(Distributions, :poissonbinomial_pdf)
-    Distributions.poissonbinomial_pdf(x::TrackedArray) = track(Distributions.poissonbinomial_pdf, x)
-    @grad function Distributions.poissonbinomial_pdf(x::TrackedArray)
-        x_data = data(x)
-        value = Distributions.poissonbinomial_pdf(x_data)
-        A = poissonbinomial_partialderivatives(x_data)
-        poissonbinomial_pdf_pullback(Δ) = (A * Δ,)
-        return value, poissonbinomial_pdf_pullback
+for f in (:poissonbinomial_pdf, :poissonbinomial_pdf_fft)
+    pullback = Symbol(f, :_pullback)
+    @eval begin
+        Distributions.$f(x::TrackedArray) = track(Distributions.$f, x)
+        @grad function Distributions.$f(x::TrackedArray)
+            x_data = data(x)
+            value = Distributions.$f(x_data)
+            A = Distributions.poissonbinomial_pdf_partialderivatives(x_data)
+            $pullback(Δ) = (A * Δ,)
+            return value, $pullback
+        end
     end
 end