test fixes

Author: mohamed82008

src/arraydist.jl

@@ -5,12 +5,22 @@ const VectorOfUnivariate = Distributions.Product
 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(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 mapvcat((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, c)
+        end)
+    end
 end
 @adjoint function Distributions.logpdf(dist::VectorOfUnivariate, x::AbstractMatrix{<:Real})
     # Any other more efficient implementation breaks Zygote
@@ -32,7 +42,9 @@ end
 function Distributions.logpdf(dist::MatrixOfUnivariate, x::AbstractMatrix{<:Real})
     # Broadcasting here breaks Tracker for some reason
     # A Zygote adjoint is defined for mapvcat to use broadcasting
-    return sum(logpdf.(dist.dists, x))
+    return sum(map(dist.dists, x) do dist, x
+        logpdf(dist, x)
+    end)
 end
 function Distributions.logpdf(dist::MatrixOfUnivariate, x::AbstractArray{<:AbstractMatrix{<:Real}})
     return mapvcat(x -> logpdf(dist, x), x)
@@ -69,7 +81,7 @@ function Distributions.logpdf(dist::VectorOfMultivariate, x::AbstractArray{<:Mat
     return mapvcat(x -> logpdf(dist, x), x)
 end
 @adjoint function Distributions.logpdf(dist::VectorOfMultivariate, x::AbstractMatrix{<:Real})
-    f(dist, x) = sum(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,9 +1,10 @@
 ## Generic ##
 
 _istracked(x) = false
+_istracked(x::TrackedArray) = false
 _istracked(x::AbstractArray{<:TrackedReal}) = true
 function mapvcat(f, args...)
-    out = f.(args...,)
+    out = map(f, args...)
     if _istracked(out)
         init = vcat(out[1])
         return reshape(reduce(vcat, drop(out, 1); init = init), size(out))
@@ -12,7 +13,7 @@ function mapvcat(f, args...)
     end
 end
 @adjoint function mapvcat(f, args...)
-    g(f, args...) = 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(logpdf.(dist, 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 = mapvcat(x -> logpdf(dist, x), x)
-        return vec(sum(reshape(temp, size(x)), dims = 1))
+        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,