Use float_type_with_fallback for logjacs and logpdfs

Author: mhauru

src/distribution_wrappers.jl

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

src/utils.jl

@@ -253,15 +253,15 @@ function (f::UnwrapSingletonTransform)(x)
 end
 
 function Bijectors.with_logabsdet_jacobian(f::UnwrapSingletonTransform, x)
-    return f(x), zero(eltype(x))
+    return f(x), zero(float_type_with_fallback(eltype(x)))
 end
 
 function Bijectors.with_logabsdet_jacobian(
     inv_f::Bijectors.Inverse{<:UnwrapSingletonTransform}, x
 )
     f = inv_f.orig
     result = reshape([x], f.input_size)
-    return result, zero(eltype(x))
+    return result, zero(float_type_with_fallback(eltype(x)))
 end
 
 """
@@ -310,22 +310,24 @@ function (inv_f::Bijectors.Inverse{<:ReshapeTransform})(x)
     return inverse(x)
 end
 
-Bijectors.with_logabsdet_jacobian(f::ReshapeTransform, x) = (f(x), zero(eltype(x)))
+function Bijectors.with_logabsdet_jacobian(f::ReshapeTransform, x)
+    return f(x), zero(float_type_with_fallback(eltype(x)))
+end
 
 function Bijectors.with_logabsdet_jacobian(inv_f::Bijectors.Inverse{<:ReshapeTransform}, x)
-    return inv_f(x), zero(eltype(x))
+    return inv_f(x), zero(float_type_with_fallback(eltype(x)))
 end
 
 struct ToChol <: Bijectors.Bijector
     uplo::Char
 end
 
 function Bijectors.with_logabsdet_jacobian(f::ToChol, x)
-    return Cholesky(Matrix(x), f.uplo, 0), zero(eltype(x))
+    return Cholesky(Matrix(x), f.uplo, 0), zero(float_type_with_fallback(eltype(x)))
 end
 
 function Bijectors.with_logabsdet_jacobian(::Bijectors.Inverse{<:ToChol}, y::Cholesky)
-    return y.UL, zero(eltype(y))
+    return y.UL, zero(float_type_with_fallback(eltype(y)))
 end
 
 function Bijectors.with_logabsdet_jacobian(::Bijectors.Inverse{<:ToChol}, y)