Merge pull request #32 from TuringLang/mt/more_tests

Many more tests and test fixes

Author: mohamed82008

.github/workflows/CI.yml

@@ -0,0 +1,27 @@
+name: CI
+
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+    types: [opened, synchronize, reopened]
+
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        julia-version: [1.0.5, 1.2.0, 1.3]
+        julia-arch: [x64, x86]
+        os: [ubuntu-latest, macOS-latest]
+        exclude:
+          - os: macOS-latest
+            julia-arch: x86
+
+    steps:
+      - uses: actions/[email protected]
+      - uses: julia-actions/setup-julia@latest
+        with:
+          version: ${{ matrix.julia-version }}
+      - uses: julia-actions/julia-runtest@master

.travis.yml

@@ -4,6 +4,7 @@ version = "0.3.2"
 
 [deps]
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
@@ -20,6 +21,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
 Combinatorics = "0.7"
+Compat = "3.6"
 DiffRules = "0.1, 1.0"
 Distributions = "0.22"
 FillArrays = "0.8"

Project.toml

@@ -8,7 +8,8 @@ using PDMats,
       Random, 
       Combinatorics,
       SpecialFunctions,
-      StatsFuns
+      StatsFuns,
+      Compat
 
 using Tracker: Tracker, TrackedReal, TrackedVector, TrackedMatrix, TrackedArray,
                 TrackedVecOrMat, track, @grad, data

src/DistributionsAD.jl

@@ -44,6 +44,12 @@ function Distributions.logpdf(dist::MatrixOfUnivariate, x::AbstractMatrix{<:Real
     # A Zygote adjoint is defined for vcatmapreduce to use broadcasting
     return sum(vcatmapreduce(logpdf, dist.dists, x))
 end
+function Distributions.logpdf(dist::MatrixOfUnivariate, x::AbstractArray{<:AbstractMatrix{<:Real}})
+    return vcatmapreduce(x -> logpdf(dist, x), x)
+end
+function Distributions.logpdf(dist::MatrixOfUnivariate, x::AbstractArray{<:Matrix{<:Real}})
+    return vcatmapreduce(x -> logpdf(dist, x), x)
+end
 function Distributions.rand(rng::Random.AbstractRNG, dist::MatrixOfUnivariate)
     return rand.(Ref(rng), dist.dists)
 end
@@ -66,6 +72,12 @@ function Distributions.logpdf(dist::VectorOfMultivariate, x::AbstractMatrix{<:Re
     # eachcol breaks Zygote, so we define an adjoint
     return sum(vcatmapreduce(logpdf, dist.dists, eachcol(x)))
 end
+function Distributions.logpdf(dist::VectorOfMultivariate, x::AbstractArray{<:AbstractMatrix{<:Real}})
+    return reshape(vcatmapreduce(x -> logpdf(dist, x), x), size(x))
+end
+function Distributions.logpdf(dist::VectorOfMultivariate, x::AbstractArray{<:Matrix{<:Real}})
+    return reshape(vcatmapreduce(x -> logpdf(dist, x), x), size(x))
+end
 @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)

src/arraydist.jl

@@ -4,8 +4,6 @@ if VERSION < v"1.1"
     eachcol(A::AbstractVecOrMat) = (view(A, :, i) for i in axes(A, 2))
 end
 
-Base.one(::Irrational) = true
-
 function vcatmapreduce(f, args...)
     init = vcat(f(first.(args)...,))
     zipped_args = zip(args...,)
@@ -14,7 +12,7 @@ function vcatmapreduce(f, args...)
     end
 end
 @adjoint function vcatmapreduce(f, args...)
-    g(f, args...) = f.(args...,)
+    g(f, args...) = f.(args...)
     return pullback(g, f, args...)
 end
 

src/common.jl

@@ -1,5 +1,7 @@
 # Univariate
 
+Tracker.dual(x::Int, p) = x
+
 const FillVectorOfUnivariate{
     S <: ValueSupport,
     T <: UnivariateDistribution{S},
@@ -46,15 +48,23 @@ end
 function _flat_logpdf(dist, x)
     if toflatten(dist)
         f, args = flatten(dist)
-        return sum(f.(args..., x))
+        if any(Tracker.istracked, args)
+            return sum(f.(args..., x))
+        else
+            return sum(logpdf.(dist, x))
+        end
     else
         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))
+        if any(Tracker.istracked, args)
+            return vec(sum(f.(args..., x), dims = 1))
+        else
+            return vec(sum(logpdf.(dist, x), dims = 1))
+        end
     else
         temp = vcatmapreduce(x -> logpdf(dist, x), x)
         return vec(sum(reshape(temp, size(x)), dims = 1))
@@ -74,7 +84,7 @@ function Distributions.logpdf(dist::FillMatrixOfUnivariate, x::AbstractMatrix{<:
     return _flat_logpdf(dist.dists.value, x)
 end
 function Distributions.rand(rng::Random.AbstractRNG, dist::FillMatrixOfUnivariate)
-    return rand(rng, dist.dists.value, length.(dist.dists.axes))
+    return rand(rng, dist.dists.value, length.(dist.dists.axes)...,)
 end
 
 # Multivariate
@@ -94,18 +104,18 @@ function Distributions.logpdf(
 )
     return _logpdf(dist, x)
 end
-@adjoint function Distributions.logpdf(
+function _logpdf(
     dist::FillVectorOfMultivariate,
     x::AbstractMatrix{<:Real},
 )
-    return pullback(_logpdf, dist, x)
+    return sum(logpdf(dist.dists.value, x))
 end
-function _logpdf(
+@adjoint function Distributions.logpdf(
     dist::FillVectorOfMultivariate,
     x::AbstractMatrix{<:Real},
 )
-    return sum(logpdf(dist.dists.value, x))
+    return pullback(_logpdf, dist, x)
 end
 function Distributions.rand(rng::Random.AbstractRNG, dist::FillVectorOfMultivariate)
-    return rand(rng, dist.dists.value, length.(dist.dists.axes))
+    return rand(rng, dist.dists.value, length.(dist.dists.axes)...,)
 end

src/filldist.jl

@@ -26,7 +26,6 @@ const flattened_dists = [   Bernoulli,
                             NegativeBinomial,
                             Poisson,
                             Skellam,
-                            PoissonBinomial,
                             Arcsine,
                             Beta,
                             BetaPrime,
@@ -42,10 +41,10 @@ const flattened_dists = [   Bernoulli,
                             FDist,
                             Frechet,
                             Gamma,
-                            GeneralizedExtremeValue,
+                            #GeneralizedExtremeValue,
                             GeneralizedPareto,
                             Gumbel,
-                            InverseGamma,
+                            #InverseGamma,
                             InverseGaussian,
                             Kolmogorov,
                             Laplace,
@@ -55,17 +54,17 @@ const flattened_dists = [   Bernoulli,
                             LogitNormal,
                             LogNormal,
                             Normal,
-                            NormalCanon,
-                            NormalInverseGaussian,
+                            #NormalCanon,
+                            #NormalInverseGaussian,
                             Pareto,
                             PGeneralizedGaussian,
                             Rayleigh,
                             SymTriangularDist,
                             TDist,
                             TriangularDist,
                             Triweight,
-                            Categorical,
-                            Truncated,
+                            #Truncated,
+                            #VonMises,
                         ]
 for T in flattened_dists
     @eval toflatten(::$T) = true

src/flatten.jl

@@ -1,3 +1,13 @@
+## MatrixBeta
+
+function Distributions.logpdf(d::MatrixBeta, X::AbstractArray{<:TrackedMatrix{<:Real}})
+    return reshape(vcatmapreduce(x -> logpdf(d, x), X), size(X))
+end
+@adjoint function Distributions.logpdf(d::MatrixBeta, X::AbstractArray{<:Matrix{<:Real}})
+    f(d, X) = map(x -> logpdf(d, x), X)
+    return pullback(f, d, X)
+end
+
 # Adapted from Distributions.jl
 
 ## Wishart
@@ -12,6 +22,13 @@ end
 
 #### Constructors
 
+function TuringWishart(d::Wishart)
+    d = TuringWishart(d.df, getchol(d.S), d.c0)
+end
+getchol(p::PDMat) = p.chol
+getchol(p::PDiagMat) = Diagonal(map(sqrt, p.diag))
+getchol(p::ScalMat) = Diagonal(fill(sqrt(p.value), p.dim))
+
 function TuringWishart(df::T, S::AbstractMatrix) where {T <: Real}
     p = size(S, 1)
     df > p - 1 || error("dpf should be greater than dim - 1.")
@@ -66,7 +83,7 @@ end
 function Distributions.entropy(d::TuringWishart)
     p = Distributions.dim(d)
     df = d.df
-    d.c0 - 0.5 * (df - p - 1) * meanlogdet(d) + 0.5 * df * p
+    d.c0 - 0.5 * (df - p - 1) * Distributions.meanlogdet(d) + 0.5 * df * p
 end
 
 #  Gupta/Nagar (1999) Theorem 3.3.15.i
@@ -82,12 +99,24 @@ end
 
 #### Evaluation
 
+function Distributions.logpdf(d::Wishart, X::TrackedMatrix)
+    return logpdf(TuringWishart(d), X)
+end
+function Distributions.logpdf(d::Wishart, X::AbstractArray{<:TrackedMatrix})
+    return logpdf(TuringWishart(d), X)
+end
 function Distributions.logpdf(d::TuringWishart, X::AbstractMatrix{<:Real})
     df = d.df
     p = Distributions.dim(d)
     Xcf = cholesky(X)
     return 0.5 * ((df - (p + 1)) * logdet(Xcf) - tr(d.chol \ X)) - d.c0
 end
+function Distributions.logpdf(d::TuringWishart, X::AbstractArray{<:AbstractMatrix{<:Real}})
+    return reshape(vcatmapreduce(x -> logpdf(d, x), X), size(X))
+end
+function Distributions.logpdf(d::TuringWishart, X::AbstractArray{<:Matrix{<:Real}})
+    return reshape(vcatmapreduce(x -> logpdf(d, x), X), size(X))
+end
 
 #### Sampling
 function Distributions._rand!(rng::AbstractRNG, d::TuringWishart, A::AbstractMatrix)
@@ -128,6 +157,13 @@ end
 
 #### Constructors
 
+function TuringInverseWishart(d::InverseWishart)
+    d = TuringInverseWishart(d.df, getmatrix(d.Ψ), d.c0)
+end
+getmatrix(p::PDMat) = p.mat
+getmatrix(p::PDiagMat) = Diagonal(p.diag)
+getmatrix(p::ScalMat) = Diagonal(fill(p.value, p.dim))
+
 function TuringInverseWishart(df::T, Ψ::AbstractMatrix) where T<:Real
     p = size(Ψ, 1)
     df > p - 1 || error("df should be greater than dim - 1.")
@@ -182,6 +218,12 @@ end
 
 #### Evaluation
 
+function Distributions.logpdf(d::InverseWishart, X::TrackedMatrix)
+    return logpdf(TuringInverseWishart(d), X)
+end
+function Distributions.logpdf(d::InverseWishart, X::AbstractArray{<:TrackedMatrix})
+    return logpdf(TuringInverseWishart(d), X)
+end
 function Distributions.logpdf(d::TuringInverseWishart, X::AbstractMatrix{<:Real})
     p = Distributions.dim(d)
     df = d.df
@@ -190,7 +232,12 @@ function Distributions.logpdf(d::TuringInverseWishart, X::AbstractMatrix{<:Real}
     Ψ = d.S
     -0.5 * ((df + p + 1) * logdet(Xcf) + tr(Xcf \ Ψ)) - d.c0
 end
-
+function Distributions.logpdf(d::TuringInverseWishart, X::AbstractArray{<:AbstractMatrix{<:Real}})
+    return reshape(vcatmapreduce(x -> logpdf(d, x), X), size(X))
+end
+function Distributions.logpdf(d::TuringInverseWishart, X::AbstractArray{<:Matrix{<:Real}})
+    return reshape(vcatmapreduce(x -> logpdf(d, x), X), size(X))
+end
 
 #### Sampling
 

src/matrixvariate.jl

@@ -116,8 +116,6 @@ struct TuringDiagMvNormal{Tm<:AbstractVector, Tσ<:AbstractVector} <: Continuous
     σ::Tσ
 end
 
-Distributions.params(d::TuringDiagMvNormal) = (d.m, d.σ)
-Distributions.dim(d::TuringDiagMvNormal) = length(d.m)
 Base.length(d::TuringDiagMvNormal) = length(d.m)
 Base.size(d::TuringDiagMvNormal) = (length(d),)
 Distributions.rand(d::TuringDiagMvNormal, n::Int...) = rand(Random.GLOBAL_RNG, d, n...)
@@ -164,6 +162,20 @@ function _logpdf(d::TuringDenseMvNormal, x::AbstractMatrix)
     return -((size(x, 1) * log(2π) + logdet(d.C)) .+ vec(sum(abs2.(zygote_ldiv(d.C.U', x .- d.m)), dims=1))) ./ 2
 end
 
+for T in (:TrackedVector, :TrackedMatrix)
+    @eval begin
+        function Distributions.logpdf(d::MvNormal{<:Any, <:PDMats.ScalMat}, x::$T)
+            logpdf(TuringScalMvNormal(d.μ, d.Σ.value), x)
+        end
+        function Distributions.logpdf(d::MvNormal{<:Any, <:PDMats.PDiagMat}, x::$T)
+            logpdf(TuringDiagMvNormal(d.μ, d.Σ.diag), x)
+        end
+        function Distributions.logpdf(d::MvNormal{<:Any, <:PDMats.PDMat}, x::$T)
+            logpdf(TuringDenseMvNormal(d.μ, d.Σ.chol), x)
+        end
+    end
+end
+
 import StatsBase: entropy
 function entropy(d::TuringDiagMvNormal)
     T = eltype(d.σ)