fix rebase and consistency issues

Author: mohamed82008

src/multivariate.jl

@@ -13,11 +13,9 @@ function TuringDenseMvNormal(m::AbstractVector, A::AbstractMatrix)
     return TuringDenseMvNormal(m, cholesky(A))
 end
 Base.length(d::TuringDenseMvNormal) = length(d.m)
-function Distributions.rand(rng::Random.AbstractRNG, d::TuringDenseMvNormal)
-    return d.m .+ d.C.U' * randn(rng, length(d))
-end
-function Distributions.rand(rng::Random.AbstractRNG, d::TuringDenseMvNormal, n::Int)
-    return d.m .+ d.C.U' * randn(rng, length(d), n)
+Distributions.rand(d::TuringDenseMvNormal, n::Int...) = rand(Random.GLOBAL_RNG, d, n...)
+function Distributions.rand(rng::Random.AbstractRNG, d::TuringDenseMvNormal, n::Int...)
+    return d.m .+ d.C.U' * randn(rng, length(d), n...)
 end
 
 """
@@ -32,43 +30,21 @@ end
 
 Base.length(d::TuringDiagMvNormal) = length(d.m)
 Base.size(d::TuringDiagMvNormal) = (length(d), length(d))
-function Distributions.rand(rng::Random.AbstractRNG, d::TuringDiagMvNormal)
-    return d.m .+ d.σ .* randn(rng, length(d))
-end
-function Distributions.rand(rng::Random.AbstractRNG, d::TuringDiagMvNormal, n::Int)
-    return d.m .+ d.σ .* randn(rng, length(d), n)
-end
-
-struct TuringScalMvNormal{Tm<:AbstractVector, Tσ<:Real} <: ContinuousMultivariateDistribution
-    m::Tm
-    σ::Tσ
-end
-
-Base.length(d::TuringScalMvNormal) = length(d.m)
-Base.size(d::TuringScalMvNormal) = (length(d), length(d))
-function Distributions.rand(rng::Random.AbstractRNG, d::TuringScalMvNormal)
-    return d.m .+ d.σ .* randn(rng, length(d))
-end
-function Distributions.rand(rng::Random.AbstractRNG, d::TuringScalMvNormal, n::Int)
-    return d.m .+ d.σ .* randn(rng, length(d), n)
-
-Base.length(d::TuringDiagMvNormal) = length(d.m)
-Base.size(d::TuringDiagMvNormal) = (length(d), length(d))
-function Distributions.rand(rng::Random.AbstractRNG, d::TuringDiagMvNormal)
-    return d.m .+ d.σ .* randn(rng, length(d))
+Distributions.rand(d::TuringDiagMvNormal, n::Int...) = rand(Random.GLOBAL_RNG, d, n...)
+function Distributions.rand(rng::Random.AbstractRNG, d::TuringDiagMvNormal, n::Int...)
+    return d.m .+ d.σ .* randn(rng, length(d), n...)
 end
 
-
 struct TuringScalMvNormal{Tm<:AbstractVector, Tσ<:Real} <: ContinuousMultivariateDistribution
     m::Tm
     σ::Tσ
 end
 
 Base.length(d::TuringScalMvNormal) = length(d.m)
 Base.size(d::TuringScalMvNormal) = (length(d), length(d))
-function Distributions.rand(rng::Random.AbstractRNG, d::TuringScalMvNormal)
-    return d.m .+ d.σ .* randn(rng, length(d))
->>>>>>> multiple distributions as one
+Distributions.rand(d::TuringScalMvNormal, n::Int...) = rand(Random.GLOBAL_RNG, d, n...)
+function Distributions.rand(rng::Random.AbstractRNG, d::TuringScalMvNormal, n::Int...)
+    return d.m .+ d.σ .* randn(rng, length(d), n...)
 end
 
 for T in (:AbstractVector, :AbstractMatrix)
@@ -95,31 +71,6 @@ function _logpdf(d::TuringDenseMvNormal, x::AbstractVector)
 end
 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
-=======
-for T in (:TrackedVector, :TrackedMatrix)
-    @eval function Distributions.logpdf(d::MvNormal{<:Any, <:PDMats.ScalMat}, x::$T)
-        logpdf(TuringScalMvNormal(d.μ, d.Σ.value), x)
-    end
-end
-
-function _logpdf(d::TuringScalMvNormal, x::AbstractVector)
-    return -(length(x) * log(2π) + 2 * sum(log(d.σ)) + sum(abs2, (x .- d.m) ./ d.σ)) / 2
-end
-function _logpdf(d::TuringScalMvNormal, x::AbstractMatrix)
-    return -(size(x, 2) * log(2π) .+ 2 * sum(log(d.σ)) .+ sum(abs2, (x .- d.m) ./ d.σ, dims=1)') ./ 2
-end
-
-function _logpdf(d::TuringDiagMvNormal, x::AbstractVector)
-    return -(length(x) * log(2π) + 2 * sum(log.(d.σ)) + sum(abs2, (x .- d.m) ./ d.σ)) / 2
-end
-function _logpdf(d::TuringDiagMvNormal, x::AbstractMatrix)
-    return -(size(x, 2) * log(2π) .+ 2 * sum(log.(d.σ)) .+ sum(abs2, (x .- d.m) ./ d.σ, dims=1)') ./ 2
-end
-function _logpdf(d::TuringDenseMvNormal, x::AbstractVector)
-    return -(length(x) * log(2π) + logdet(d.C) + sum(abs2, zygote_ldiv(d.C.U', x .- d.m))) / 2
-end
-function _logpdf(d::TuringDenseMvNormal, x::AbstractMatrix)
-    return -(size(x, 2) * log(2π) .+ logdet(d.C) .+ sum(abs2, zygote_ldiv(d.C.U', x .- d.m), dims=1)') ./ 2
 end
 
 # zero mean, dense covariance
@@ -129,9 +80,9 @@ MvNormal(A::TrackedMatrix) = TuringMvNormal(A)
 MvNormal(σ::TrackedVector) = TuringMvNormal(σ)
 
 # dense mean, dense covariance
-MvNormal(m::TrackedVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringDenseMvNormal(m, A)
-MvNormal(m::TrackedVector{<:Real}, A::Matrix{<:Real}) = TuringDenseMvNormal(m, A)
-MvNormal(m::AbstractVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringDenseMvNormal(m, A)
+MvNormal(m::TrackedVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringMvNormal(m, A)
+MvNormal(m::TrackedVector{<:Real}, A::Matrix{<:Real}) = TuringMvNormal(m, A)
+MvNormal(m::AbstractVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringMvNormal(m, A)
 
 # dense mean, diagonal covariance
 function MvNormal(
@@ -237,9 +188,9 @@ MvLogNormal(A::TrackedMatrix) = TuringMvLogNormal(TuringMvNormal(A))
 MvLogNormal(σ::TrackedVector) = TuringMvLogNormal(TuringMvNormal(σ))
 
 # dense mean, dense covariance
-MvLogNormal(m::TrackedVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringMvLogNormal(TuringDenseMvNormal(m, A))
-MvLogNormal(m::TrackedVector{<:Real}, A::Matrix{<:Real}) = TuringMvLogNormal(TuringDenseMvNormal(m, A))
-MvLogNormal(m::AbstractVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringMvLogNormal(TuringDenseMvNormal(m, A))
+MvLogNormal(m::TrackedVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, A))
+MvLogNormal(m::TrackedVector{<:Real}, A::Matrix{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, A))
+MvLogNormal(m::AbstractVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, A))
 
 # dense mean, diagonal covariance
 function MvLogNormal(

src/univariate.jl

@@ -33,39 +33,25 @@ uniformlogpdf(a::TrackedReal, b::TrackedReal, x::TrackedReal) = track(uniformlog
 Tracker.@grad function uniformlogpdf(a, b, x)
     diff = data(b) - data(a)
     T = typeof(diff)
-    l = -log(diff)
-    f = isfinite(l)
-    da = 1/diff
-    n = T(NaN)
-    return l, Δ->(f ? da : n, f ? -da : n, f ? zero(T) : n)
+    if a <= data(x) <= b && a < b
+        l = -log(diff)
+        da = 1/diff^2
+        return l, Δ -> (da * Δ, -da * Δ, zero(T) * Δ)
+    else
+        n = T(NaN)
+        return l, Δ -> (n, n, n)
+    end
 end
 ZygoteRules.@adjoint function uniformlogpdf(a, b, x)
     diff = b - a
     T = typeof(diff)
-    l = -log(diff)
-    f = isfinite(l)
-    da = 1/diff
-    n = T(NaN)
-    z = zero(T)
-    return l, Δ -> (f ? (z, z, z) : (n, n, n))
-end
-for T in (:TrackedReal, :Real)
-    @eval @grad function uniformlogpdf(
-        a::TrackedReal,
-        b::TrackedReal,
-        x::$T,
-    )
-        ad = data(a)
-        bd = data(b)
-        T = typeof(a)
-        l = logpdf(Uniform(ad, bd), x)
-        f = isfinite(l)
-        temp = 1/(bd - ad)^2
-        dlda = temp
-        dldb = -temp
+    if a <= x <= b && a < b
+        l = -log(diff)
+        da = 1/diff^2
+        return l, Δ -> (da * Δ, -da * Δ, zero(T) * Δ)
+    else
         n = T(NaN)
-        z = zero(T)
-        return l, Δ -> (f ? (dlda * Δ, dldb * Δ, z) : (n, n, n))
+        return l, Δ -> (n, n, n)
     end
 end
 ZygoteRules.@adjoint function Distributions.Uniform(args...)
@@ -159,9 +145,9 @@ M, f, arity = DiffRules.@define_diffrule DistributionsAD.semicirclelogpdf(r, x)
 da, db = DiffRules.diffrule(M, f, :a, :b)
 f = :($M.$f)
 @eval begin
-    @grad $f(a::TrackedReal, b::TrackedReal) = $f(data(a), data(b)), Δ -> (Δ * $da, Δ * $db)
-    @grad $f(a::TrackedReal, b::Real) = $f(data(a), b), Δ -> (Δ * $da, Tracker._zero(b))
-    @grad $f(a::Real, b::TrackedReal) = $f(a, data(b)), Δ -> (Tracker._zero(a), Δ * $db)
+    Tracker.@grad $f(a::TrackedReal, b::TrackedReal) = $f(data(a), data(b)), Δ -> (Δ * $da, Δ * $db)
+    Tracker.@grad $f(a::TrackedReal, b::Real) = $f(data(a), b), Δ -> (Δ * $da, Tracker._zero(b))
+    Tracker.@grad $f(a::Real, b::TrackedReal) = $f(a, data(b)), Δ -> (Tracker._zero(a), Δ * $db)
     $f(a::TrackedReal, b::TrackedReal)  = track($f, a, b)
     $f(a::TrackedReal, b::Real) = track($f, a, b)
     $f(a::Real, b::TrackedReal) = track($f, a, b)
@@ -292,6 +278,7 @@ ZygoteRules.@adjoint function poissonbinomial_pdf_fft(x::AbstractArray)
         ((ForwardDiff.jacobian(x -> poissonbinomial_pdf_fft(x), x)::Matrix{T})' * Δ,)
     end
 end
+
 # The code below doesn't work because of bugs in Zygote. The above is inefficient.
 #=
 ZygoteRules.@adjoint function poissonbinomial_pdf_fft(x::AbstractArray{<:Real})