fix David's comments

Author: mohamed82008

src/DistributionsAD.jl

@@ -12,6 +12,7 @@ using PDMats,
 
 using Tracker: Tracker, TrackedReal, TrackedVector, TrackedMatrix, TrackedArray,
                 TrackedVecOrMat, track, @grad, data
+using SpecialFunctions: logabsgamma, digamma
 using ZygoteRules: ZygoteRules, @adjoint, pullback
 using LinearAlgebra: copytri!
 using Distributions: AbstractMvLogNormal, 
@@ -39,15 +40,15 @@ export TuringScalMvNormal,
        TuringPoissonBinomial,
        TuringWishart,
        TuringInverseWishart,
-       ArrayDist,
-       FillDist
+       arraydist,
+       filldist
 
 include("common.jl")
 include("univariate.jl")
 include("multivariate.jl")
 include("matrixvariate.jl")
 include("flatten.jl")
-include("array_dist.jl")
-include("multi.jl")
+include("arraydist.jl")
+include("filldist.jl")
 
 end

src/array_dist.jl renamed to src/arraydist.jl

@@ -1,43 +1,31 @@
 # Univariate
 
-const VectorOfUnivariate{
-    S <: ValueSupport,
-    Tdist <: UnivariateDistribution{S},
-    Tdists <: AbstractVector{Tdist},
-} = Distributions.Product{S, Tdist, Tdists}
-
-function ArrayDist(dists::AbstractVector{<:Normal{T}}) where {T}
-    if T <: TrackedReal
-        init_m = vcat(dists[1].μ)
-        means = mapreduce(vcat, drop(dists, 1); init = init_m) do d
-            d.μ
-        end
-        init_v = vcat(dists[1].σ^2)
-        vars = mapreduce(vcat, drop(dists, 1); init = init_v) do d
-            d.σ^2
-        end
-    else
-        means = [d.μ for d in dists]
-        vars = [d.σ^2 for d in dists]
-    end
+const VectorOfUnivariate = Distributions.Product
 
+function arraydist(dists::AbstractVector{<:Normal{T}}) where {T}
+    means = mean.(dists)
+    vars = var.(dists)
     return MvNormal(means, vars)
 end
-function ArrayDist(dists::AbstractVector{<:UnivariateDistribution})
-    return Distributions.Product(dists)
+function arraydist(dists::AbstractVector{<:Normal{<:TrackedReal}})
+    means = vcatmapreduce(mean, dists)
+    vars = vcatmapreduce(var, dists)
+    return MvNormal(means, vars)
+end
+function arraydist(dists::AbstractVector{<:UnivariateDistribution})
+    return product_distribution(dists)
 end
 function Distributions.logpdf(dist::VectorOfUnivariate, x::AbstractVector{<:Real})
-    return sum(logpdf.(dist.v, x))
+    return sum(vcatmapreduce(logpdf, dist.v, x))
 end
 function Distributions.logpdf(dist::VectorOfUnivariate, x::AbstractMatrix{<:Real})
-    # Any other more efficient implementation breaks Zygote
-    return [logpdf(dist, x[:,i]) for i in 1:size(x, 2)]
+    # eachcol breaks Zygote, so we need an adjoint
+    return vcatmapreduce((dist, c) -> logpdf.(dist, c), dist.v, eachcol(x))
 end
-function Distributions.logpdf(
-    dist::VectorOfUnivariate,
-    x::AbstractVector{<:AbstractMatrix{<:Real}},
-)
-    return logpdf.(Ref(dist), x)
+@adjoint function Distributions.logpdf(dist::VectorOfUnivariate, x::AbstractMatrix{<:Real})
+    # Any other more efficient implementation breaks Zygote
+    f(dist, x) = [sum(logpdf.(dist.v, view(x, :, i))) for i in 1:size(x, 2)]
+    return pullback(f, dist, x)
 end
 
 struct MatrixOfUnivariate{
@@ -48,14 +36,13 @@ struct MatrixOfUnivariate{
     dists::Tdists
 end
 Base.size(dist::MatrixOfUnivariate) = size(dist.dists)
-function ArrayDist(dists::AbstractMatrix{<:UnivariateDistribution})
+function arraydist(dists::AbstractMatrix{<:UnivariateDistribution})
     return MatrixOfUnivariate(dists)
 end
 function Distributions.logpdf(dist::MatrixOfUnivariate, x::AbstractMatrix{<:Real})
     # Broadcasting here breaks Tracker for some reason
-    return sum(zip(dist.dists, x)) do (dist, x)
-        logpdf(dist, x)
-    end
+    # A Zygote adjoint is defined for vcatmapreduce to use broadcasting
+    return sum(vcatmapreduce(logpdf, dist.dists, x))
 end
 function Distributions.rand(rng::Random.AbstractRNG, dist::MatrixOfUnivariate)
     return rand.(Ref(rng), dist.dists)
@@ -72,17 +59,16 @@ struct VectorOfMultivariate{
 end
 Base.size(dist::VectorOfMultivariate) = (length(dist.dists[1]), length(dist))
 Base.length(dist::VectorOfMultivariate) = length(dist.dists)
-function ArrayDist(dists::AbstractVector{<:MultivariateDistribution})
+function arraydist(dists::AbstractVector{<:MultivariateDistribution})
     return VectorOfMultivariate(dists)
 end
 function Distributions.logpdf(dist::VectorOfMultivariate, x::AbstractMatrix{<:Real})
-    return sum(logpdf(dist.dists[i], x[:,i]) for i in 1:length(dist))
+    # eachcol breaks Zygote, so we define an adjoint
+    return sum(vcatmapreduce(logpdf, dist.dists, eachcol(x)))
 end
-function Distributions.logpdf(
-    dist::VectorOfMultivariate,
-    x::AbstractVector{<:AbstractVector{<:Real}},
-)
-    return sum(logpdf(dist.dists[i], x[i]) for i in 1:length(dist))
+@adjoint function Distributions.logpdf(dist::VectorOfMultivariate, x::AbstractMatrix{<:Real})
+    f(dist, x) = sum(vcatmapreduce(i -> logpdf(dist.dists[i], view(x, :, i)), 1:size(x, 2)))
+    return pullback(f, dist, x)
 end
 function Distributions.rand(rng::Random.AbstractRNG, dist::VectorOfMultivariate)
     init = reshape(rand(rng, dist.dists[1]), :, 1)

src/common.jl

@@ -2,6 +2,18 @@
 
 Base.one(::Irrational) = 1
 
+function vcatmapreduce(f, args...)
+    init = vcat(f(first.(args)...,))
+    zipped_args = zip(args...,)
+    return mapreduce(vcat, drop(zipped_args, 1); init = init) do zarg
+        f(zarg...,)
+    end
+end
+@adjoint function vcatmapreduce(f, args...)
+    g(f, args...) = f.(args...,)
+    return pullback(g, f, args...)
+end
+
 function Base.fill(
     value::TrackedReal,
     dims::Vararg{Union{Integer, AbstractUnitRange}},
@@ -110,9 +122,12 @@ end
 
 # SpecialFunctions
 
-function SpecialFunctions.logabsgamma(x::TrackedReal)
-    v = loggamma(x)
-    return v, sign(data(v))
+SpecialFunctions.logabsgamma(x::TrackedReal) = track(logabsgamma, x)
+@grad function SpecialFunctions.logabsgamma(x::Real)
+    return logabsgamma(data(x)), Δ -> (digamma(data(x)) * Δ[1],)
+end
+@adjoint function SpecialFunctions.logabsgamma(x::Real)
+    return logabsgamma(x), Δ -> (digamma(x) * Δ[1],)
 end
 
 # Some Tracker fixes

src/multi.jl renamed to src/filldist.jl

@@ -6,36 +6,58 @@ const FillVectorOfUnivariate{
     Tdists <: Fill{T, 1},
 } = VectorOfUnivariate{S, T, Tdists}
 
-function FillDist(dist::UnivariateDistribution, N::Int)
-    return Product(Fill(dist, N))
+function filldist(dist::UnivariateDistribution, N::Int)
+    return product_distribution(Fill(dist, N))
 end
-FillDist(d::Normal, N::Int) = MvNormal(fill(d.μ, N), d.σ)
+filldist(d::Normal, N::Int) = MvNormal(fill(d.μ, N), d.σ)
+
 function Distributions.logpdf(
     dist::FillVectorOfUnivariate,
     x::AbstractVector{<:Real},
 )
-    return _flat_logpdf(dist.v.value, x)
+    return _logpdf(dist, x)
 end
 function Distributions.logpdf(
     dist::FillVectorOfUnivariate,
     x::AbstractMatrix{<:Real},
+)
+    return _logpdf(dist, x)
+end
+@adjoint function Distributions.logpdf(
+    dist::FillVectorOfUnivariate,
+    x::AbstractMatrix{<:Real},
+)
+    return pullback(_logpdf, dist, x)
+end
+
+function _logpdf(
+    dist::FillVectorOfUnivariate,
+    x::AbstractVector{<:Real},
+)
+    return _flat_logpdf(dist.v.value, x)
+end
+function _logpdf(
+    dist::FillVectorOfUnivariate,
+    x::AbstractMatrix{<:Real},
 )
     return _flat_logpdf_mat(dist.v.value, x)
 end
+
 function _flat_logpdf(dist, x)
     if toflatten(dist)
         f, args = flatten(dist)
         return sum(f.(args..., x))
     else
-        return sum(logpdf.(dist, x))
+        return sum(vcatmapreduce(x -> logpdf(dist, x), x))
     end
 end
 function _flat_logpdf_mat(dist, x)
     if toflatten(dist)
         f, args = flatten(dist)
         return vec(sum(f.(args..., x), dims = 1))
     else
-        return vec(sum(logpdf.(dist, x), dims = 1))
+        temp = vcatmapreduce(x -> logpdf(dist, x), x)
+        return vec(sum(reshape(temp, size(x)), dims = 1))
     end
 end
 
@@ -45,7 +67,7 @@ const FillMatrixOfUnivariate{
     Tdists <: Fill{T, 2},
 } = MatrixOfUnivariate{S, T, Tdists}
 
-function FillDist(dist::UnivariateDistribution, N1::Integer, N2::Integer)
+function filldist(dist::UnivariateDistribution, N1::Integer, N2::Integer)
     return MatrixOfUnivariate(Fill(dist, N1, N2))
 end
 function Distributions.logpdf(dist::FillMatrixOfUnivariate, x::AbstractMatrix{<:Real})
@@ -63,12 +85,24 @@ const FillVectorOfMultivariate{
     Tdists <: Fill{T, 1},
 } = VectorOfMultivariate{S, T, Tdists}
 
-function FillDist(dist::MultivariateDistribution, N::Int)
+function filldist(dist::MultivariateDistribution, N::Int)
     return VectorOfMultivariate(Fill(dist, N))
 end
 function Distributions.logpdf(
     dist::FillVectorOfMultivariate,
     x::AbstractMatrix{<:Real},
+)
+    return _logpdf(dist, x)
+end
+@adjoint function Distributions.logpdf(
+    dist::FillVectorOfMultivariate,
+    x::AbstractMatrix{<:Real},
+)
+    return pullback(_logpdf, dist, x)
+end
+function _logpdf(
+    dist::FillVectorOfMultivariate,
+    x::AbstractMatrix{<:Real},
 )
     return sum(logpdf(dist.dists.value, x))
 end

src/multivariate.jl

@@ -59,7 +59,7 @@ function simplex_logpdf(alpha, lmnB, x::AbstractVector)
     sum((alpha .- 1) .* log.(x)) - lmnB
 end
 function simplex_logpdf(alpha, lmnB, x::AbstractMatrix)
-    @views init = vcat(sum((alpha .- 1) .* log.(x[:,1])))
+    init = vcat(sum((alpha .- 1) .* log.(view(x, :, 1))))
     mapreduce(vcat, drop(eachcol(x), 1); init = init) do c
         sum((alpha .- 1) .* log.(c)) - lmnB
     end

test/distributions.jl

@@ -159,8 +159,8 @@ separator()
 @testset "Multivariate discrete distributions" begin
     test_head("Testing: Multivariate discrete distributions")
     mult_disc_dists = [
-        DistSpec(:((p) -> FillDist(Bernoulli(p), dim)), (0.45,), fill(1, dim)),
-        DistSpec(:((p) -> ArrayDist(fill(Bernoulli(p), dim))), (0.45,), fill(1, dim)),
+        DistSpec(:((p) -> filldist(Bernoulli(p), dim)), (0.45,), fill(1, dim)),
+        DistSpec(:((p) -> arraydist(fill(Bernoulli(p), dim))), (0.45,), fill(1, dim)),
         DistSpec(:((p) -> Multinomial(2, p / sum(p))), (fill(0.5, 2),), [2, 0]),
     ]
     for d in mult_disc_dists
@@ -176,10 +176,10 @@ separator()
     test_head("Testing: Multivariate continuous distributions")
     mult_cont_dists = [
         # Vector case
-        DistSpec(:(() -> FillDist(Beta(), dim)), (), fill(0.5, dim)),
-        DistSpec(:(() -> ArrayDist(fill(Beta(), dim))), (), fill(0.5, dim)),
-        DistSpec(:((m, v) -> FillDist(Normal(m, sqrt(v)), dim)), (1.0, 1.0), norm_val_vec),
-        DistSpec(:((m, v) -> ArrayDist(fill(Normal(m, sqrt(v)), dim))), (1.0, 1.0), norm_val_vec),
+        DistSpec(:(() -> filldist(Beta(), dim)), (), fill(0.5, dim)),
+        DistSpec(:(() -> arraydist(fill(Beta(), dim))), (), fill(0.5, dim)),
+        DistSpec(:((m, v) -> filldist(Normal(m, sqrt(v)), dim)), (1.0, 1.0), norm_val_vec),
+        DistSpec(:((m, v) -> arraydist(fill(Normal(m, sqrt(v)), dim))), (1.0, 1.0), norm_val_vec),
         DistSpec(:MvNormal, (mean, cov_mat), norm_val_vec),
         DistSpec(:MvNormal, (mean, cov_vec), norm_val_vec),
         DistSpec(:MvNormal, (mean, Diagonal(cov_vec)), norm_val_vec),
@@ -198,10 +198,10 @@ separator()
         DistSpec(:MvLogNormal, (Diagonal(cov_vec),), norm_val_vec),
         DistSpec(:(cov_num -> MvLogNormal(dim, cov_num)), (cov_num,), norm_val_vec),
         # Matrix case
-        DistSpec(:(() -> FillDist(Beta(), dim)), (), fill(0.5, dim, dim)),
-        DistSpec(:(() -> ArrayDist(fill(Beta(), dim))), (), fill(0.5, dim, dim)),
-        DistSpec(:((m, v) -> FillDist(Normal(m, sqrt(v)), dim)), (1.0, 1.0), norm_val_mat),
-        DistSpec(:((m, v) -> ArrayDist(fill(Normal(m, sqrt(v)), dim))), (1.0, 1.0), norm_val_mat),
+        DistSpec(:(() -> filldist(Beta(), dim)), (), fill(0.5, dim, dim)),
+        DistSpec(:(() -> arraydist(fill(Beta(), dim))), (), fill(0.5, dim, dim)),
+        DistSpec(:((m, v) -> filldist(Normal(m, sqrt(v)), dim)), (1.0, 1.0), norm_val_mat),
+        DistSpec(:((m, v) -> arraydist(fill(Normal(m, sqrt(v)), dim))), (1.0, 1.0), norm_val_mat),
         DistSpec(:MvNormal, (mean, cov_vec), norm_val_mat),
         DistSpec(:MvNormal, (mean, Diagonal(cov_vec)), norm_val_mat),
         DistSpec(:MvNormal, (mean, cov_num), norm_val_mat),
@@ -254,10 +254,10 @@ separator()
 @testset "Matrix-variate continuous distributions" begin
     test_head("Testing: Matrix-variate continuous distributions")
     matrix_cont_dists = [
-        DistSpec(:(() -> FillDist(Beta(), dim, dim)), (), fill(0.5, dim, dim)),
-        DistSpec(:(() -> ArrayDist(fill(Beta(), dim, dim))), (), fill(0.5, dim, dim)),
-        DistSpec(:((m, v) -> FillDist(Normal(m, sqrt(v)), dim, 2)), (1.0, 1.0), norm_val_mat),
-        DistSpec(:((m, v) -> ArrayDist(fill(Normal(m, sqrt(v)), dim, 2))), (1.0, 1.0), norm_val_mat),
+        DistSpec(:(() -> filldist(Beta(), dim, dim)), (), fill(0.5, dim, dim)),
+        DistSpec(:(() -> arraydist(fill(Beta(), dim, dim))), (), fill(0.5, dim, dim)),
+        DistSpec(:((m, v) -> filldist(Normal(m, sqrt(v)), dim, 2)), (1.0, 1.0), norm_val_mat),
+        DistSpec(:((m, v) -> arraydist(fill(Normal(m, sqrt(v)), dim, 2))), (1.0, 1.0), norm_val_mat),
         DistSpec(:((n1, n2)->MatrixBeta(dim, n1, n2)), (dim, dim), beta_mat),
         DistSpec(:Wishart, (dim, cov_mat), cov_mat),
         DistSpec(:InverseWishart, (dim, cov_mat), cov_mat),