Fix zero-type of logjac for ReshapeTransform (#851)

* Fix zero-type of logjac for ReshapeTransform

* Bump patch version to 0.35.4

* Make logpdf of NoDist be of the eltype of the argument

* Use float_type_with_fallback for logjacs and logpdfs

* Make LogProbType be float(Real)

Author: mhauru

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.35.3"
+version = "0.35.4"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/distribution_wrappers.jl

@@ -54,30 +54,38 @@ function Distributions.rand!(
 ) where {N}
     return Distributions.rand!(rng, d.dist, x)
 end
-Distributions.logpdf(d::NoDist{<:Univariate}, ::Real) = 0
-Distributions.logpdf(d::NoDist{<:Multivariate}, ::AbstractVector{<:Real}) = 0
-function Distributions.logpdf(d::NoDist{<:Multivariate}, x::AbstractMatrix{<:Real})
-    return zeros(Int, size(x, 2))
+function Distributions.logpdf(::NoDist{<:Univariate}, x::Real)
+    return zero(LogProbType)
+end
+function Distributions.logpdf(::NoDist{<:Multivariate}, x::AbstractVector{<:Real})
+    return zero(LogProbType)
+end
+function Distributions.logpdf(::NoDist{<:Multivariate}, x::AbstractMatrix{<:Real})
+    return zeros(LogProbType, size(x, 2))
+end
+function Distributions.logpdf(::NoDist{<:Matrixvariate}, x::AbstractMatrix{<:Real})
+    return zero(LogProbType)
 end
-Distributions.logpdf(d::NoDist{<:Matrixvariate}, ::AbstractMatrix{<:Real}) = 0
 Distributions.minimum(d::NoDist) = minimum(d.dist)
 Distributions.maximum(d::NoDist) = maximum(d.dist)
 
-Bijectors.logpdf_with_trans(d::NoDist{<:Univariate}, ::Real, ::Bool) = 0
+function Bijectors.logpdf_with_trans(::NoDist{<:Univariate}, x::Real, ::Bool)
+    return zero(LogProbType)
+end
 function Bijectors.logpdf_with_trans(
-    d::NoDist{<:Multivariate}, ::AbstractVector{<:Real}, ::Bool
+    ::NoDist{<:Multivariate}, x::AbstractVector{<:Real}, ::Bool
 )
-    return 0
+    return zero(LogProbType)
 end
 function Bijectors.logpdf_with_trans(
-    d::NoDist{<:Multivariate}, x::AbstractMatrix{<:Real}, ::Bool
+    ::NoDist{<:Multivariate}, x::AbstractMatrix{<:Real}, ::Bool
 )
-    return zeros(Int, size(x, 2))
+    return zeros(LogProbType, size(x, 2))
 end
 function Bijectors.logpdf_with_trans(
-    d::NoDist{<:Matrixvariate}, ::AbstractMatrix{<:Real}, ::Bool
+    ::NoDist{<:Matrixvariate}, x::AbstractMatrix{<:Real}, ::Bool
 )
-    return 0
+    return zero(LogProbType)
 end
 
 Bijectors.bijector(d::NoDist) = Bijectors.bijector(d.dist)

src/utils.jl

@@ -5,6 +5,16 @@ const NO_DEFAULT = NoDefault()
 # A short-hand for a type commonly used in type signatures for VarInfo methods.
 VarNameTuple = NTuple{N,VarName} where {N}
 
+# TODO(mhauru) This is currently used in the transformation functions of NoDist,
+# ReshapeTransform, and UnwrapSingletonTransform, and in VarInfo. We should also use it in
+# SimpleVarInfo and maybe other places.
+"""
+The type for all log probability variables.
+
+This is Float64 on 64-bit systems and Float32 on 32-bit systems.
+"""
+const LogProbType = float(Real)
+
 """
     @addlogprob!(ex)
 
@@ -252,12 +262,16 @@ function (f::UnwrapSingletonTransform)(x)
     return only(x)
 end
 
-Bijectors.with_logabsdet_jacobian(f::UnwrapSingletonTransform, x) = (f(x), 0)
+function Bijectors.with_logabsdet_jacobian(f::UnwrapSingletonTransform, x)
+    return f(x), zero(LogProbType)
+end
+
 function Bijectors.with_logabsdet_jacobian(
     inv_f::Bijectors.Inverse{<:UnwrapSingletonTransform}, x
 )
     f = inv_f.orig
-    return (reshape([x], f.input_size), 0)
+    result = reshape([x], f.input_size)
+    return result, zero(LogProbType)
 end
 
 """
@@ -306,18 +320,26 @@ function (inv_f::Bijectors.Inverse{<:ReshapeTransform})(x)
     return inverse(x)
 end
 
-Bijectors.with_logabsdet_jacobian(f::ReshapeTransform, x) = (f(x), 0)
+function Bijectors.with_logabsdet_jacobian(f::ReshapeTransform, x)
+    return f(x), zero(LogProbType)
+end
 
 function Bijectors.with_logabsdet_jacobian(inv_f::Bijectors.Inverse{<:ReshapeTransform}, x)
-    return (inv_f(x), 0)
+    return inv_f(x), zero(LogProbType)
 end
 
 struct ToChol <: Bijectors.Bijector
     uplo::Char
 end
 
-Bijectors.with_logabsdet_jacobian(f::ToChol, x) = (Cholesky(Matrix(x), f.uplo, 0), 0)
-Bijectors.with_logabsdet_jacobian(::Bijectors.Inverse{<:ToChol}, y::Cholesky) = (y.UL, 0)
+function Bijectors.with_logabsdet_jacobian(f::ToChol, x)
+    return Cholesky(Matrix(x), f.uplo, 0), zero(LogProbType)
+end
+
+function Bijectors.with_logabsdet_jacobian(::Bijectors.Inverse{<:ToChol}, y::Cholesky)
+    return y.UL, zero(LogProbType)
+end
+
 function Bijectors.with_logabsdet_jacobian(::Bijectors.Inverse{<:ToChol}, y)
     return error(
         "Inverse{ToChol} is only defined for Cholesky factorizations. " *

src/varinfo.jl

@@ -821,7 +821,7 @@ end
 
 # VarInfo
 
-VarInfo(meta=Metadata()) = VarInfo(meta, Ref{Float64}(0.0), Ref(0))
+VarInfo(meta=Metadata()) = VarInfo(meta, Ref{LogProbType}(0.0), Ref(0))
 
 function TypedVarInfo(vi::VectorVarInfo)
     new_metas = group_by_symbol(vi.metadata)