BatchMvNormal variance shape (#32)

* let BatchMvNormal always return variance of shape (xlength,batchsize)

Project.toml

@@ -1,7 +1,7 @@
 name = "ConditionalDists"
 uuid = "c648c4dd-c1e0-49a6-84b9-144ae7fd2468"
 authors = ["Niklas Heim <[email protected]>"]
-version = "0.4.1"
+version = "0.4.2"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/batch_mvnormal.jl

@@ -17,10 +17,8 @@ BatchMvNormal(μ::AbstractMatrix{T}, σ::AbstractMatrix{T}) where T<:Real = Batc
 Base.eltype(d::BMN) = eltype(d.μ)
 Distributions.params(d::BMN) = (d.μ, d.σ)
 Distributions.mean(d::BMN) = d.μ
-Distributions.var(d::BMN) = d.σ .^2
-
-#Distributions.var(d::BatchScalMvNormal) = fill(similar(d.σ,size(d.μ,1)),1) .* reshape(d.σ .^2,1,:)
-#Distributions.var(d::BatchScalMvNormal) = ones(eltype(d), size(d.μ,1), 1) .* reshape(d.σ .^2, 1, :)
+Distributions.var(d::BatchDiagMvNormal) = d.σ .^2
+Distributions.var(d::BatchScalMvNormal) = fillsimilar(d.σ,size(d.μ,1),1) .* reshape(d.σ .^2,1,:)
 
 function Distributions.rand(d::BatchDiagMvNormal)
     μ, σ = d.μ, d.σ
@@ -34,16 +32,9 @@ function Distributions.rand(d::BatchScalMvNormal)
     μ .+ σ .* r
 end
 
-function Distributions.logpdf(d::BatchDiagMvNormal, x::AbstractMatrix{T}) where T<:Real
+function Distributions.logpdf(d::BMN, x::AbstractMatrix{T}) where T<:Real
     n = size(d.μ,1)
     μ = mean(d)
     σ2 = var(d)
     -(vec(sum(((x - μ).^2) ./ σ2 .+ log.(σ2), dims=1)) .+ n*log(T(2π))) / 2
 end
-
-function Distributions.logpdf(d::BatchScalMvNormal, x::AbstractMatrix{T}) where T<:Real
-    n = size(d.μ,1)
-    μ = mean(d)
-    σ2 = reshape(var(d), 1, :)
-    -(vec(sum(((x - μ).^2) ./ σ2 .+ log.(σ2), dims=1)) .+ n*log(T(2π))) / 2
-end

test/cond_mvnormal.jl

@@ -54,7 +54,7 @@
     σ2 = var(res)
     @test res isa ConditionalDists.BatchScalMvNormal
     @test size(μ) == (xlength,batchsize)
-    @test size(σ2) == (batchsize,)
+    @test size(σ2) == (xlength,batchsize)
 
     x = rand(Float32, xlength, batchsize) |> gpu
     z = rand(Float32, zlength, batchsize) |> gpu
@@ -78,7 +78,7 @@
     σ2 = var(res)
     @test res isa ConditionalDists.BatchScalMvNormal
     @test size(μ) == (xlength,batchsize)
-    @test size(σ2) == (batchsize,)
+    @test size(σ2) == (xlength,batchsize)
 
     x = rand(Float32, xlength, batchsize) |> gpu
     z = rand(Float32, zlength, batchsize) |> gpu