Merge pull request #29 from TuringLang/mt/categorical_dirichlet

Categorical and Dirichlet

Author: mohamed82008

src/multivariate.jl

@@ -1,3 +1,91 @@
+## Dirichlet ##
+
+struct TuringDirichlet{T, TV <: AbstractVector} <: ContinuousMultivariateDistribution
+    alpha::TV
+    alpha0::T
+    lmnB::T
+end
+function check(alpha)
+    all(ai -> ai > 0, alpha) || 
+        throw(ArgumentError("Dirichlet: alpha must be a positive vector."))
+end
+Zygote.@nograd DistributionsAD.check
+
+function TuringDirichlet(alpha::AbstractVector)
+    check(alpha)
+    alpha0 = sum(alpha)
+    lmnB = sum(loggamma, alpha) - loggamma(alpha0)
+    T = promote_type(typeof(alpha0), typeof(lmnB))
+    TV = typeof(alpha)
+    TuringDirichlet{T, TV}(alpha, alpha0, lmnB)
+end
+
+function TuringDirichlet(d::Integer, alpha::Real)
+    alpha0 = alpha * d
+    _alpha = fill(alpha, d)
+    lmnB = loggamma(alpha) * d - loggamma(alpha0)
+    T = promote_type(typeof(alpha0), typeof(lmnB))
+    TV = typeof(_alpha)
+    TuringDirichlet{T, TV}(_alpha, alpha0, lmnB)
+end
+function TuringDirichlet(alpha::AbstractVector{T}) where {T <: Integer}
+    Tf = float(T)
+    TuringDirichlet(convert(AbstractVector{Tf}, alpha))
+end
+TuringDirichlet(d::Integer, alpha::Integer) = TuringDirichlet(d, Float64(alpha))
+
+Distributions.Dirichlet(alpha::TrackedVector) = TuringDirichlet(alpha)
+Distributions.Dirichlet(d::Integer, alpha::TrackedReal) = TuringDirichlet(d, alpha)
+
+function Distributions.logpdf(d::TuringDirichlet, x::AbstractVector)
+    simplex_logpdf(d.alpha, d.lmnB, x)
+end
+function Distributions.logpdf(d::TuringDirichlet, x::AbstractMatrix)
+    simplex_logpdf(d.alpha, d.lmnB, x)
+end
+function Distributions.logpdf(d::Dirichlet{T}, x::TrackedVecOrMat) where {T}
+    TV = typeof(d.alpha)
+    logpdf(TuringDirichlet{T, TV}(d.alpha, d.alpha0, d.lmnB), x)
+end
+
+ZygoteRules.@adjoint function Distributions.Dirichlet(alpha)
+    return pullback(TuringDirichlet, alpha)
+end
+ZygoteRules.@adjoint function Distributions.Dirichlet(d, alpha)
+    return pullback(TuringDirichlet, d, alpha)
+end
+
+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])) - lmnB)
+    mapreduce(vcat, drop(eachcol(x), 1); init = init) do c
+        sum((alpha .- 1) .* log.(c)) - lmnB
+    end
+end
+
+Tracker.@grad function simplex_logpdf(alpha, lmnB, x::AbstractVector)
+    simplex_logpdf(data(alpha), data(lmnB), data(x)), Δ -> begin
+        (Δ .* log.(data(x)), -Δ, Δ .* (data(alpha) .- 1))
+    end
+end
+Tracker.@grad function simplex_logpdf(alpha, lmnB, x::AbstractMatrix)
+    simplex_logpdf(data(alpha), data(lmnB), data(x)), Δ -> begin
+        (log.(data(x)) * Δ, -sum(Δ), repeat(data(alpha) .- 1, 1, size(x, 2)) * Diagonal(Δ))
+    end
+end
+
+ZygoteRules.@adjoint function simplex_logpdf(alpha, lmnB, x::AbstractVector)
+    simplex_logpdf(alpha, lmnB, x), Δ -> (Δ .* log.(x), -Δ, Δ .* (alpha .- 1))
+end
+
+ZygoteRules.@adjoint function simplex_logpdf(alpha, lmnB, x::AbstractMatrix)
+    simplex_logpdf(alpha, lmnB, x), Δ -> begin
+        (log.(x) * Δ, -sum(Δ), repeat(alpha .- 1, 1, size(x, 2)) * Diagonal(Δ))
+    end
+end
+
 ## MvNormal ##
 
 """
@@ -68,6 +156,7 @@ end
 function _logpdf(d::TuringDiagMvNormal, x::AbstractMatrix)
     return -((size(x, 1) * log(2π) + 2 * sum(log.(d.σ))) .+ vec(sum(abs2.((x .- d.m) ./ d.σ), dims=1))) ./ 2
 end
+
 function _logpdf(d::TuringDenseMvNormal, x::AbstractVector)
     return -(length(x) * log(2π) + logdet(d.C) + sum(abs2.(zygote_ldiv(d.C.U', x .- d.m)))) / 2
 end

src/univariate.jl

@@ -337,4 +337,14 @@ function _dft_zygote(x::Vector{T}) where T
     end
     return copy(y)
 end
-=#
+=#
+
+## Categorical ##
+
+function Base.convert(
+    ::Type{Distributions.DiscreteNonParametric{T,P,Ts,Ps}},
+    d::Distributions.DiscreteNonParametric{T,P,Ts,Ps},
+) where {T<:Real,P<:Real,Ts<:AbstractVector{T},Ps<:AbstractVector{P}}
+    DiscreteNonParametric{T,P,Ts,Ps}(support(d), probs(d), check_args=false)
+end
+

test/distributions.jl

@@ -215,6 +215,7 @@ separator()
         DistSpec(:MvLogNormal, (cov_vec,), norm_val_mat),
         DistSpec(:MvLogNormal, (Diagonal(cov_vec),), norm_val_mat),
         DistSpec(:(cov_num -> MvLogNormal(dim, cov_num)), (cov_num,), norm_val_mat),
+        DistSpec(:Dirichlet, (alpha,), dir_val),
     ]
 
     broken_mult_cont_dists = [
@@ -231,7 +232,6 @@ separator()
         DistSpec(:MvNormalCanon, (cov_mat,), norm_val_mat),
         DistSpec(:MvNormalCanon, (cov_vec,), norm_val_mat),
         DistSpec(:(cov_num -> MvNormalCanon(dim, cov_num)), (cov_num,), norm_val_mat),
-        DistSpec(:Dirichlet, (alpha,), dir_val),
         # Test failure
         DistSpec(:MvNormal, (mean, cov_mat), norm_val_mat),
         DistSpec(:MvNormal, (cov_mat,), norm_val_mat),