Introduce strict_logdensityof

Author: oschulz

ext/MeasureBaseChainRulesCoreExt.jl

@@ -44,11 +44,16 @@ ChainRulesCore.rrule(::typeof(checked_arg), ν, x) = checked_arg(ν, x), _checke
 
 # = return type inference ====================================================
 
-using MeasureBase: logdensityof_rt
+using MeasureBase: logdensityof_rt, strict_logdensityof_rt
 
 _logdensityof_rt_pullback(::Any) = (NoTangent(), NoTangent(), ZeroTangent())
 function ChainRulesCore.rrule(::typeof(logdensityof_rt), target, v)
     logdensityof_rt(target, v), _logdensityof_rt_pullback
 end
 
+_strict_logdensityof_rt_pullback(::Any) = (NoTangent(), NoTangent(), ZeroTangent())
+function ChainRulesCore.rrule(::typeof(strict_logdensityof_rt), target, v)
+    strict_logdensityof_rt(target, v), _strict_logdensityof_rt_pullback
+end
+
 end # module MeasureBaseChainRulesCoreExt

src/combinators/half.jl

@@ -19,8 +19,8 @@ function Base.rand(rng::AbstractRNG, ::Type{T}, μ::Half) where {T}
     return abs(rand(rng, T, unhalf(μ)))
 end
 
-function logdensityof(μ::Half, x)
-    ld = logdensityof(unhalf(μ), x) - loghalf
+function strict_logdensityof(μ::Half, x)
+    ld = strict_logdensityof(unhalf(μ), x) - loghalf
     return x ≥ 0 ? ld : oftype(ld, -Inf)
 end
 

src/combinators/power.jl

@@ -78,7 +78,7 @@ params(d::PowerMeasure) = params(first(marginals(d)))
     basemeasure(d.parent)^d.axes
 end
 
-for func in [:logdensityof, :logdensity_def]
+for func in [:strict_logdensityof, :logdensity_def]
     @eval @inline function $func(d::PowerMeasure{M}, x) where {M}
         parent = d.parent
         sum(x) do xj

src/combinators/product.jl

@@ -72,7 +72,7 @@ function _rand_product(
     end |> collect
 end
 
-for func in [:logdensityof, :logdensity_def]
+for func in [:strict_logdensityof, :logdensity_def]
     @eval @inline function $func(d::AbstractProductMeasure, x)
         mapreduce($func, +, marginals(d), x)
     end
@@ -82,7 +82,7 @@ struct ProductMeasure{M} <: AbstractProductMeasure
     marginals::M
 end
 
-@inline function logdensity_rel(μ::ProductMeasure, ν::ProductMeasure, x)
+@inline function strict_logdensity_rel(μ::ProductMeasure, ν::ProductMeasure, x)
     mapreduce(logdensity_rel, +, marginals(μ), marginals(ν), x)
 end
 
@@ -109,7 +109,7 @@ end
     return q
 end
 
-for func in [:logdensityof, :logdensity_def]
+for func in [:strict_logdensityof, :logdensity_def]
     # For tuples, `mapreduce` has trouble with type inference
     @eval @inline function $func(d::ProductMeasure{T}, x) where {T<:Tuple}
         ℓs = map($func, marginals(d), x)

src/combinators/spikemixture.jl

@@ -21,7 +21,7 @@ end
     SpikeMixture(basemeasure(μ.m), static(1.0), static(1.0))
 end
 
-for func in [:logdensityof, :logdensity_def]
+for func in [:strict_logdensityof, :logdensity_def]
     @eval @inline function $func(μ::SpikeMixture, x)
         # NOTE: We could instead write this as 
         # R1 = typeof(log(one(μ.s))) 

src/combinators/transformedmeasure.jl

@@ -103,15 +103,15 @@ function Pretty.tile(ν::PushforwardMeasure)
 end
 
 # TODO: THIS IS ALMOST CERTAINLY WRONG 
-# @inline function logdensity_rel(
+# @inline function strict_logdensity_rel(
 #     ν::PushforwardMeasure{FF1,IF1,M1,<:AdaptRootMeasure},
 #     β::PushforwardMeasure{FF2,IF2,M2,<:AdaptRootMeasure},
 #     y,
 # ) where {FF1,IF1,M1,FF2,IF2,M2}
 #     x = β.inv_f(y)
 #     f = ν.inv_f ∘ β.f
 #     inv_f = β.inv_f ∘ ν.f
-#     logdensity_rel(pushfwd(f, inv_f, ν.origin, AdaptRootMeasure()), β.origin, x)
+#     strict_logdensity_rel(pushfwd(f, inv_f, ν.origin, AdaptRootMeasure()), β.origin, x)
 # end
 
 # TODO: Would profit from custom pullback:
@@ -132,7 +132,7 @@ function _combine_logd_with_ladj(logd_orig::Real, ladj::Real)
     end
 end
 
-function logdensityof(
+function strict_logdensityof(
     @nospecialize(μ::_NonBijectivePusfwdMeasure{M,<:PushfwdRootMeasure}),
     @nospecialize(v::Any)
 ) where {M}
@@ -143,7 +143,7 @@ function logdensityof(
     )
 end
 
-function logdensityof(
+function strict_logdensityof(
     @nospecialize(μ::_NonBijectivePusfwdMeasure{M,<:AdaptRootMeasure}),
     @nospecialize(v::Any)
 ) where {M}
@@ -154,7 +154,7 @@ function logdensityof(
     )
 end
 
-for func in [:logdensityof, :logdensity_def]
+for func in [:strict_logdensityof, :logdensity_def]
     @eval function $func(ν::PushforwardMeasure{F,I,M,<:AdaptRootMeasure}, y) where {F,I,M}
         f_inv = unwrap(ν.finv)
         x, inv_ladj = with_logabsdet_jacobian(f_inv, y)

src/density-core.jl

@@ -26,11 +26,15 @@ To compute log-density relative to `basemeasure(m)` or *define* a log-density
 `logdensity_def`. 
 
 To compute a log-density relative to a specific base-measure, see
-`logdensity_rel`. 
+`logdensity_rel`.
+
+# Implementation
+
+Do not specialize `logdensityof` directly for subtypes of `AbstractMeasure`,
+specialize `MeasureBase.logdensity_def` and `MeasureBase.strict_logdensityof` instead.
 """
-@inline function logdensityof(μ::AbstractMeasure, x)
-    result = dynamic(unsafe_logdensityof(μ, x))
-    _checksupport(insupport(μ, x), result)
+@inline function logdensityof(μ::AbstractMeasure, x) #!!!!!!!!!!!!!!!!!
+    strict_logdensityof(μ, x)
 end
 
 @inline function logdensityof_rt(::T, ::U) where {T,U}
@@ -41,6 +45,24 @@ _checksupport(cond, result) = ifelse(cond == true, result, oftype(result, -Inf))
 
 export unsafe_logdensityof
 
+"""
+    MeasureBase.strict_logdensityof(μ, x)
+
+Compute the log-density of the measure `μ` at `x` relative to `rootmeasure(m)`.
+In contrast to [`logdensityof(μ, x)`](@ref), this will not take implicit pushforwards
+of `μ` (depending on the type of `x`) into account.
+"""
+function strict_logdensityof end
+
+@inline function strict_logdensityof(μ, x)
+    result = dynamic(unsafe_logdensityof(μ, x))
+    _checksupport(insupport(μ, x), result)
+end
+
+@inline function strict_logdensityof_rt(::T, ::U) where {T,U}
+    Core.Compiler.return_type(strict_logdensityof, Tuple{T,U})
+end
+
 # https://discourse.julialang.org/t/counting-iterations-to-a-type-fixpoint/75876/10?u=cscherrer
 """
     unsafe_logdensityof(m, x)
@@ -68,14 +90,27 @@ See also `logdensityof`.
 end
 
 """
-    logdensity_rel(m1, m2, x)
+    logdensity_rel(μ, ν, x)
 
-Compute the log-density of `m1` relative to `m2` at `x`. This function checks
-whether `x` is in the support of `m1` or `m2` (or both, or neither). If `x` is
+Compute the log-density of `μ` relative to `ν` at `x`. This function checks
+whether `x` is in the support of `μ` or `ν` (or both, or neither). If `x` is
 known to be in the support of both, it can be more efficient to call
-`unsafe_logdensity_rel`. 
+`unsafe_logdensity_rel`.
+"""
+function logdensity_rel(μ, ν, x)
+    strict_logdensity_rel(μ, ν, x)
+end
+
 """
-@inline function logdensity_rel(μ::M, ν::N, x::X) where {M,N,X}
+    MeasureBase.strict_logdensity_rel(μ, ν, x)
+
+Compute the log-density of `μ` relative to `ν` at `x`. In contrast to
+[`logdensity_rel(μ, ν, x)`](@ref), this will not take implicit pushforwards
+of `μ` and `ν` (depending on the type of `x`) into account.
+"""
+function strict_logdensity_rel end
+
+@inline function strict_logdensity_rel(μ::M, ν::N, x::X) where {M,N,X}
     T = unstatic(
         promote_type(
             return_type(logdensity_def, (μ, x)),

src/density.jl

@@ -163,10 +163,10 @@ logdensity_def(μ::DensityMeasure, x) = logdensityof(μ.f, x)
 
 density_def(μ::DensityMeasure, x) = densityof(μ.f, x)
 
-function logdensityof(μ::DensityMeasure, x::Any)
+function strict_logdensityof(μ::DensityMeasure, x::Any)
     integrand, μ_base = μ.f, μ.base
 
-    base_logval = logdensityof(μ_base, x)
+    base_logval = strict_logdensityof(μ_base, x)
 
     T = typeof(base_logval)
     U = logdensityof_rt(integrand, x)

src/primitive.jl

@@ -19,8 +19,8 @@ basemeasure(μ::PrimitiveMeasure) = μ
 
 @inline basemeasure_depth(::PrimitiveMeasure) = static(0)
 
-@inline logdensityof(::PrimitiveMeasure, x::Real) = zero(float(typeof(x)))
-@inline logdensityof(::PrimitiveMeasure, x) = false
+@inline strict_logdensityof(::PrimitiveMeasure, x::Real) = zero(float(typeof(x)))
+@inline strict_logdensityof(::PrimitiveMeasure, x) = false
 
 logdensity_def(::PrimitiveMeasure, x) = static(0.0)
 

src/primitives/counting.jl

@@ -12,14 +12,14 @@ struct Counting{T} <: AbstractMeasure
     Counting(supp) = new{Core.Typeof(supp)}(supp)
 end
 
-@inline function logdensityof(μ::Counting, x::Real)
+@inline function strict_logdensityof(μ::Counting, x::Real)
     R = float(typeof(x))
     insupport(μ, x) ? zero(R) : R(-Inf)
 end
 
-@inline logdensityof(μ::Counting, x) = insupport(μ, x) ? 0.0 : -Inf
+@inline strict_logdensityof(μ::Counting, x) = insupport(μ, x) ? 0.0 : -Inf
 
-@inline logdensity_def(μ::Counting, x) = logdensityof(μ, x)
+@inline logdensity_def(μ::Counting, x) = strict_logdensityof(μ, x)
 
 basemeasure(::Counting) = CountingBase()