Merge pull request #37 from TuringLang/mt/perf_fixes

Minor performance and bug fixes (lessons learnt from TuringExamples)

Author: mohamed82008

.github/workflows/ForwardDiff_Tracker.yml

@@ -0,0 +1,29 @@
+name: ForwardDiff and Tracker tests
+
+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
+        env:
+          STAGE: ForwardDiff_Tracker

.github/workflows/CI.yml renamed to .github/workflows/Others.yml

@@ -1,4 +1,4 @@
-name: CI
+name: Other tests
 
 on:
   push:
@@ -25,3 +25,5 @@ jobs:
         with:
           version: ${{ matrix.julia-version }}
       - uses: julia-actions/julia-runtest@master
+        env:
+          STAGE: Others

.github/workflows/Zygote.yml

@@ -0,0 +1,29 @@
+name: Zygote tests
+
+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
+        env:
+          STAGE: Zygote

src/arraydist.jl

@@ -2,25 +2,25 @@
 
 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{<: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 arraydist(dists::AbstractVector{<:Normal})
+    m = mapvcat(mean, dists)
+    v = mapvcat(var, dists)
+    return MvNormal(m, v)
+end
+
 function Distributions.logpdf(dist::VectorOfUnivariate, x::AbstractVector{<:Real})
-    return sum(vcatmapreduce(logpdf, dist.v, x))
+    return sum(map((d, x) -> logpdf(d, x), dist.v, x))
 end
 function Distributions.logpdf(dist::VectorOfUnivariate, x::AbstractMatrix{<:Real})
     # eachcol breaks Zygote, so we need an adjoint
-    return vcatmapreduce((dist, c) -> logpdf.(dist, c), dist.v, eachcol(x))
+    return mapvcat(dist.v, eachcol(x)) do dist, c
+        sum(map(c) do x
+            logpdf(dist, x)
+        end)
+    end
 end
 @adjoint function Distributions.logpdf(dist::VectorOfUnivariate, x::AbstractMatrix{<:Real})
     # Any other more efficient implementation breaks Zygote
@@ -41,14 +41,16 @@ function arraydist(dists::AbstractMatrix{<:UnivariateDistribution})
 end
 function Distributions.logpdf(dist::MatrixOfUnivariate, x::AbstractMatrix{<:Real})
     # Broadcasting here breaks Tracker for some reason
-    # A Zygote adjoint is defined for vcatmapreduce to use broadcasting
-    return sum(vcatmapreduce(logpdf, dist.dists, x))
+    # A Zygote adjoint is defined for mapvcat to use broadcasting
+    return sum(map(dist.dists, x) do dist, x
+        logpdf(dist, x)
+    end)
 end
 function Distributions.logpdf(dist::MatrixOfUnivariate, x::AbstractArray{<:AbstractMatrix{<:Real}})
-    return vcatmapreduce(x -> logpdf(dist, x), x)
+    return mapvcat(x -> logpdf(dist, x), x)
 end
 function Distributions.logpdf(dist::MatrixOfUnivariate, x::AbstractArray{<:Matrix{<:Real}})
-    return vcatmapreduce(x -> logpdf(dist, x), x)
+    return mapvcat(x -> logpdf(dist, x), x)
 end
 function Distributions.rand(rng::Random.AbstractRNG, dist::MatrixOfUnivariate)
     return rand.(Ref(rng), dist.dists)
@@ -70,16 +72,16 @@ function arraydist(dists::AbstractVector{<:MultivariateDistribution})
 end
 function Distributions.logpdf(dist::VectorOfMultivariate, x::AbstractMatrix{<:Real})
     # eachcol breaks Zygote, so we define an adjoint
-    return sum(vcatmapreduce(logpdf, dist.dists, eachcol(x)))
+    return sum(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))
+    return mapvcat(x -> logpdf(dist, x), x)
 end
 function Distributions.logpdf(dist::VectorOfMultivariate, x::AbstractArray{<:Matrix{<:Real}})
-    return reshape(vcatmapreduce(x -> logpdf(dist, x), x), size(x))
+    return mapvcat(x -> logpdf(dist, x), 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)))
+    f(dist, x) = sum(mapvcat(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)

src/common.jl

@@ -1,14 +1,19 @@
 ## Generic ##
 
-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...,)
+_istracked(x) = false
+_istracked(x::TrackedArray) = false
+_istracked(x::AbstractArray{<:TrackedReal}) = true
+function mapvcat(f, args...)
+    out = map(f, args...)
+    if _istracked(out)
+        init = vcat(out[1])
+        return reshape(reduce(vcat, drop(out, 1); init = init), size(out))
+    else
+        return out
     end
 end
-@adjoint function vcatmapreduce(f, args...)
-    g(f, args...) = f.(args...)
+@adjoint function mapvcat(f, args...)
+    g(f, args...) = map(f, args...)
     return pullback(g, f, args...)
 end
 

src/filldist.jl

@@ -48,26 +48,20 @@ end
 function _flat_logpdf(dist, x)
     if toflatten(dist)
         f, args = flatten(dist)
-        if any(Tracker.istracked, args)
-            return sum(f.(args..., x))
-        else
-            return sum(logpdf.(dist, x))
-        end
+        return sum(f.(args..., x))
     else
-        return sum(vcatmapreduce(x -> logpdf(dist, x), x))
+        return sum(mapvcat(x) do x
+            logpdf(dist, x)
+        end)
     end
 end
 function _flat_logpdf_mat(dist, x)
     if toflatten(dist)
         f, args = flatten(dist)
-        if any(Tracker.istracked, args)
-            return vec(sum(f.(args..., x), dims = 1))
-        else
-            return vec(sum(logpdf.(dist, x), dims = 1))
-        end
+        return vec(sum(f.(args..., x), dims = 1))
     else
-        temp = vcatmapreduce(x -> logpdf(dist, x), x)
-        return vec(sum(reshape(temp, size(x)), dims = 1))
+        temp = mapvcat(x -> logpdf(dist, x), x)
+        return vec(sum(temp, dims = 1))
     end
 end
 

src/flatten.jl

@@ -41,7 +41,7 @@ const flattened_dists = [   Bernoulli,
                             FDist,
                             Frechet,
                             Gamma,
-                            #GeneralizedExtremeValue,
+                            GeneralizedExtremeValue,
                             GeneralizedPareto,
                             Gumbel,
                             #InverseGamma,
@@ -63,6 +63,7 @@ const flattened_dists = [   Bernoulli,
                             TDist,
                             TriangularDist,
                             Triweight,
+                            TuringUniform,
                             #Truncated,
                             #VonMises,
                         ]

src/matrixvariate.jl

@@ -1,7 +1,7 @@
 ## MatrixBeta
 
 function Distributions.logpdf(d::MatrixBeta, X::AbstractArray{<:TrackedMatrix{<:Real}})
-    return reshape(vcatmapreduce(x -> logpdf(d, x), X), size(X))
+    return mapvcat(x -> logpdf(d, x), X)
 end
 @adjoint function Distributions.logpdf(d::MatrixBeta, X::AbstractArray{<:Matrix{<:Real}})
     f(d, X) = map(x -> logpdf(d, x), X)
@@ -112,10 +112,10 @@ function Distributions.logpdf(d::TuringWishart, X::AbstractMatrix{<:Real})
     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))
+    return mapvcat(x -> logpdf(d, x), X)
 end
 function Distributions.logpdf(d::TuringWishart, X::AbstractArray{<:Matrix{<:Real}})
-    return reshape(vcatmapreduce(x -> logpdf(d, x), X), size(X))
+    return mapvcat(x -> logpdf(d, x), X)
 end
 
 #### Sampling
@@ -233,10 +233,10 @@ function Distributions.logpdf(d::TuringInverseWishart, X::AbstractMatrix{<:Real}
     -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))
+    return mapvcat(x -> logpdf(d, x), X)
 end
 function Distributions.logpdf(d::TuringInverseWishart, X::AbstractArray{<:Matrix{<:Real}})
-    return reshape(vcatmapreduce(x -> logpdf(d, x), X), size(X))
+    return mapvcat(x -> logpdf(d, x), X)
 end
 
 #### Sampling

src/univariate.jl

@@ -13,6 +13,8 @@ function TuringUniform(a::Real, b::Real)
     return TuringUniform{T}(T(a), T(b))
 end
 Distributions.logpdf(d::TuringUniform, x::Real) = uniformlogpdf(d.a, d.b, x)
+Base.minimum(d::TuringUniform) = d.a
+Base.maximum(d::TuringUniform) = d.b
 
 Distributions.Uniform(a::TrackedReal, b::Real) = TuringUniform{TrackedReal}(a, b)
 Distributions.Uniform(a::Real, b::TrackedReal) = TuringUniform{TrackedReal}(a, b)
@@ -348,3 +350,21 @@ function Base.convert(
     DiscreteNonParametric{T,P,Ts,Ps}(support(d), probs(d), check_args=false)
 end
 
+# Fix SubArray support
+function Distributions.DiscreteNonParametric{T,P,Ts,Ps}(
+    vs::Ts,
+    ps::Ps;
+    check_args=true,
+) where {T<:Real, P<:Real, Ts<:AbstractVector{T}, Ps<:SubArray{P, 1}}
+    cps = ps[:]
+    return DiscreteNonParametric{T,P,Ts,typeof(cps)}(vs, cps; check_args = check_args)
+end
+
+function Distributions.DiscreteNonParametric{T,P,Ts,Ps}(
+    vs::Ts,
+    ps::Ps;
+    check_args=true,
+) where {T<:Real,P<:Real,Ts<:AbstractVector{T},Ps<:TrackedArray{P, 1, <:SubArray{P, 1}}}
+    cps = ps[:]
+    return DiscreteNonParametric{T,P,Ts,typeof(cps)}(vs, cps; check_args = check_args)
+end