perf fixes

Author: mohamed82008

src/arraydist.jl

@@ -2,25 +2,15 @@
 
 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 Distributions.logpdf(dist::VectorOfUnivariate, x::AbstractVector{<:Real})
-    return sum(vcatmapreduce(logpdf, dist.v, x))
+    return sum(logpdf.(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, c) -> logpdf.(dist, c), dist.v, eachcol(x))
 end
 @adjoint function Distributions.logpdf(dist::VectorOfUnivariate, x::AbstractMatrix{<:Real})
     # Any other more efficient implementation breaks Zygote
@@ -41,14 +31,14 @@ 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(logpdf.(dist.dists, x))
 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 +60,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(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,13 +1,17 @@
 ## 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::AbstractArray{<:TrackedReal}) = true
+function mapvcat(f, args...)
+    out = 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...)
+@adjoint function mapvcat(f, args...)
     g(f, args...) = f.(args...)
     return pullback(g, f, args...)
 end

src/filldist.jl

@@ -54,7 +54,7 @@ function _flat_logpdf(dist, x)
             return sum(logpdf.(dist, x))
         end
     else
-        return sum(vcatmapreduce(x -> logpdf(dist, x), x))
+        return sum(logpdf.(dist, x))
     end
 end
 function _flat_logpdf_mat(dist, x)
@@ -66,7 +66,7 @@ function _flat_logpdf_mat(dist, x)
             return vec(sum(logpdf.(dist, x), dims = 1))
         end
     else
-        temp = vcatmapreduce(x -> logpdf(dist, x), x)
+        temp = mapvcat(x -> logpdf(dist, x), x)
         return vec(sum(reshape(temp, size(x)), dims = 1))
     end
 end

src/flatten.jl

@@ -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