WIP: Major upgrade

Project.toml

@@ -1,16 +1,20 @@
 name = "MeasureBase"
 uuid = "fa1605e6-acd5-459c-a1e6-7e635759db14"
-version = "0.14.12"
+version = "0.15.0"
 authors = ["Chad Scherrer <chad.scherrer@gmail.com>", "Oliver Schulz <oschulz@mpp.mpg.de>", "contributors"]
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConstantRNGs = "aa9b60e7-6b1c-4c29-a6e5-e43521412437"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 FunctionChains = "8e6b2b91-af83-483e-ba35-d00930e4cf9b"
+HeterogeneousComputing = "2182be2a-124f-4a91-8389-f06db5907a21"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
@@ -20,32 +24,62 @@ LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899"
 MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+OneTwoMany = "762dc654-8631-413a-a342-372a7419ad9d"
 PrettyPrinting = "54e16d92-306c-5ea0-a30b-337be88ac337"
 PropertyFunctions = "09e99361-2bb8-48a2-a80f-de58f0739eb4"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StaticThings = "7e4b4f32-fbf9-4b74-9510-4d15222ac973"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 
 [weakdeps]
+AffineMaps = "2c83c9a8-abf5-4329-a0d7-deffaf474661"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 
 [extensions]
+MeasureBaseAffineMapsExt = "AffineMaps"
 MeasureBaseChainRulesCoreExt = "ChainRulesCore"
+MeasureBaseDistributionsExt = ["Distributions", "StatsBase", "StatsFuns", "PDMats"]
+MeasureBaseDistributionsChainRulesCoreExt = ["Distributions", "ChainRulesCore"]
+MeasureBaseDistributionsForwardDiffExt = ["Distributions", "ForwardDiff"]
+MeasureBaseDistributionsForwardDiffPullbacksExt = ["Distributions", "ForwardDiffPullbacks", "ChainRulesCore"]
+MeasureBaseDistributionsMooncakeExt = ["Distributions", "Mooncake"]
+MeasureBaseForwardDiffExt = "ForwardDiff"
+MeasureBaseForwardDiffPullbacksExt = "ForwardDiffPullbacks"
+MeasureBaseMooncakeExt = "Mooncake"
+MeasureBaseReactantExt = "Reactant"
 
 [compat]
+Adapt = "3.7, 4"
+AffineMaps = "0.3"
+ArgCheck = "1, 2"
+ArraysOfArrays = "1.3"
 ChainRulesCore = "1"
 ChangesOfVariables = "0.1.3"
 Compat = "3.35, 4"
 ConstantRNGs = "0.1.1"
 ConstructionBase = "1.3"
 DensityInterface = "0.4"
+Distributions = "0.25.1"
 FillArrays = "0.12, 0.13, 1"
-FunctionChains = "0.2"
+ForwardDiff = "0.10, 1"
+ForwardDiffPullbacks = "0.2"
+FunctionChains = "0.2.3"
+HeterogeneousComputing = "0.2.3"
 IfElse = "0.1"
 IntervalSets = "0.7"
 InverseFunctions = "0.1.8"
@@ -54,15 +88,23 @@ LinearAlgebra = "1"
 LogExpFunctions = "0.3, 1"
 LogarithmicNumbers = "1"
 MappedArrays = "0.4"
+Mooncake = "0.5.34"
 NaNMath = "0.3, 1"
+OneTwoMany = "0.1.2"
+PDMats = "0.11"
 PrettyPrinting = "0.3, 0.4"
-PropertyFunctions = "0.2.2"
+PropertyFunctions = "0.3"
 Random = "1"
+Reactant = "0.2"
 Reexport = "1"
-SpecialFunctions = "2"
+SpecialFunctions = "2.1.4"
 Static = "0.8, 1"
 StaticArrays = "1.5"
+StaticThings = "0.3"
 Statistics = "1"
+StructArrays = "0.6.18, 0.7"
+StatsBase = "0.33, 0.34"
+StatsFuns = "0.9, 1, 2"
 Test = "1"
 Tricks = "0.1"
 julia = "1.10"

docs/make.jl

@@ -12,7 +12,7 @@ DocMeta.setdocmeta!(MeasureBase, :DocTestSetup, :(using MeasureBase); recursive
 
 makedocs(
     sitename = "MeasureBase",
-    modules = [MeasureBase],
+    modules = [MeasureBase, MeasureBase.MeasureOperators],
     format = Documenter.HTML(
         prettyurls = !("local" in ARGS),
         canonical = "https://juliamath.github.io/MeasureBase.jl/stable/",

docs/src/api.md

@@ -21,6 +21,6 @@ Order = [:macro, :function]
 # Documentation
 
 ```@autodocs
-Modules = [MeasureBase]
+Modules = [MeasureBase, MeasureBase.MeasureOperators]
 Order = [:module, :type, :constant, :macro, :function]
 ```

ext/MeasureBaseAffineMapsExt.jl

@@ -0,0 +1,45 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseAffineMapsExt
+
+using MeasureBase
+using MeasureBase: PushforwardMeasure, AdaptRootMeasure, PushfwdRootMeasure
+using MeasureBase: StaticInteger
+using AffineMaps: AbstractAffineMap
+using ChangesOfVariables: with_logabsdet_jacobian
+
+# Affine maps treat matrices as batches of column vectors, so flat batches
+# of vector variates are applied as `(n, :)` matrices:
+function MeasureBase._apply_generic(f::AbstractAffineMap, X::AbstractArray, ::StaticInteger{1})
+    _columns_back(f(_as_columns(X)), X)
+end
+
+@inline _as_columns(x::AbstractVector) = x
+@inline _as_columns(X::AbstractArray) = reshape(X, (size(X, 1), :))
+@inline _columns_back(y::AbstractVector, ::AbstractVector) = y
+@inline _columns_back(Y::AbstractMatrix, X::AbstractArray) = reshape(Y, (size(Y, 1), Base.tail(size(X))...))
+
+const _AffinePushfwd{M,S} = PushforwardMeasure{<:AbstractAffineMap,<:AbstractAffineMap,M,S}
+
+# Densities of affine pushforwards of vector variates use the per-column
+# log-abs-det-Jacobians of the inverse map:
+for (bhead, head) in [(:batched_logdensityof_impl, :logdensityof_impl), (:batched_logdensity_def, :logdensity_def)]
+    @eval function MeasureBase.$bhead(ν::_AffinePushfwd{M,<:AdaptRootMeasure}, Y::AbstractArray) where {M}
+        _affine_pushfwd_ld(MeasureBase.$head, ν, Y, MeasureBase._static_ndims(ν))
+    end
+    @eval function MeasureBase.$bhead(ν::_AffinePushfwd{M,<:PushfwdRootMeasure}, Y::AbstractArray) where {M}
+        MeasureBase._batched_kernel(MeasureBase.$head, ν.origin, MeasureBase._apply_batched(ν.finv, Y, MeasureBase._static_ndims(ν)))
+    end
+end
+
+function _affine_pushfwd_ld(f::F, ν::PushforwardMeasure, Y::AbstractArray, ::StaticInteger{1}) where {F}
+    X2, ladj2 = with_logabsdet_jacobian(ν.finv, _as_columns(Y))
+    ℓ = MeasureBase._batched_kernel(f, ν.origin, _columns_back(X2, Y))
+    return MeasureBase._lazy_combine_ladj(ℓ, _ladj_back(ladj2, Y))
+end
+_affine_pushfwd_ld(f::F, ν::PushforwardMeasure, Y::AbstractArray, k) where {F} = MeasureBase._default_batched_kernel(f, ν, Y, k)
+
+@inline _ladj_back(ladj::Number, ::AbstractVector) = ladj
+@inline _ladj_back(ladj::AbstractMatrix, Y::AbstractArray) = reshape(ladj, Base.tail(size(Y)))
+
+end # module MeasureBaseAffineMapsExt

ext/MeasureBaseChainRulesCoreExt.jl

@@ -3,22 +3,84 @@
 module MeasureBaseChainRulesCoreExt
 
 using MeasureBase
-using ChainRulesCore: NoTangent, ZeroTangent
+using ChainRulesCore: NoTangent, ZeroTangent, @non_differentiable
 import ChainRulesCore
 
 # = utils ====================================================================
 
 using MeasureBase: isneginf, isposinf
 
 _isneginf_pullback(::Any) = (NoTangent(), ZeroTangent())
-ChainRulesCore.rrule(::typeof(isneginf), x) = isneginf(x), _logdensityof_rt_pullback
+ChainRulesCore.rrule(::typeof(isneginf), x) = isneginf(x), _isneginf_pullback
 
 _isposinf_pullback(::Any) = (NoTangent(), ZeroTangent())
 ChainRulesCore.rrule(::typeof(isposinf), x) = isposinf(x), _isposinf_pullback
 
+using MeasureBase: _adignore_call
+
+@inline _adignore_call_pullback(@nospecialize ΔΩ) = (NoTangent(), NoTangent())
+ChainRulesCore.rrule(::typeof(_adignore_call), f) = _adignore_call(f), _adignore_call_pullback
+
+using MeasureBase: convert_realtype
+
+_convert_realtype_pullback(ΔΩ) = NoTangent(), NoTangent(), ΔΩ
+ChainRulesCore.rrule(::typeof(convert_realtype), ::Type{T}, x) where {T} =
+    convert_realtype(T, x), _convert_realtype_pullback
+
+# = collection utils =========================================================
+
+using MeasureBase: _pushfront, _pushback, _rev_cumsum, _exp_cumsum_log
+
+function ChainRulesCore.rrule(::typeof(_pushfront), v::AbstractVector, x)
+    result = _pushfront(v, x)
+    function _pushfront_pullback(thunked_ΔΩ)
+        ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ)
+        (NoTangent(), ΔΩ[firstindex(ΔΩ)+1:lastindex(ΔΩ)], ΔΩ[firstindex(ΔΩ)])
+    end
+    return result, _pushfront_pullback
+end
+
+
+function ChainRulesCore.rrule(::typeof(_pushback), v::AbstractVector, x)
+    result = _pushback(v, x)
+    function _pushback_pullback(thunked_ΔΩ)
+        ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ)
+        (NoTangent(), ΔΩ[firstindex(ΔΩ):lastindex(ΔΩ)-1], ΔΩ[lastindex(ΔΩ)])
+    end
+    return result, _pushback_pullback
+end
+
+
+function ChainRulesCore.rrule(::typeof(_rev_cumsum), xs::AbstractVector)
+    result = _rev_cumsum(xs)
+    function _rev_cumsum_pullback(ΔΩ)
+        ∂xs = ChainRulesCore.@thunk cumsum(ChainRulesCore.unthunk(ΔΩ))
+        (NoTangent(), ∂xs)
+    end
+    return result, _rev_cumsum_pullback
+end
+
+
+function ChainRulesCore.rrule(::typeof(_exp_cumsum_log), xs::AbstractVector)
+    result = _exp_cumsum_log(xs)
+    function _exp_cumsum_log_pullback(ΔΩ)
+        ∂xs = inv.(xs) .* _rev_cumsum(exp.(cumsum(log.(xs))) .* ChainRulesCore.unthunk(ΔΩ))
+        (NoTangent(), ∂xs)
+    end
+    return result, _exp_cumsum_log_pullback
+end
+
+
 # = insupport & friends ======================================================
 
-using MeasureBase: check_dof, require_insupport, checked_arg, _checksupport, _origin_depth
+# Variate layouts are type-level information:
+using MeasureBase: fixed_stream_size, _pushfwd_varsize, _unit_bounds, _prob_floor
+@non_differentiable fixed_stream_size(::Type)
+@non_differentiable _pushfwd_varsize(f, μ)
+@non_differentiable _unit_bounds(p)
+@non_differentiable _prob_floor(p)
+
+using MeasureBase: check_dof, require_insupport, checked_arg, _checksupport
 
 @inline function ChainRulesCore.rrule(::typeof(_checksupport), cond, result)
     y = _checksupport(cond, result)
@@ -33,22 +95,10 @@ function ChainRulesCore.rrule(::typeof(require_insupport), μ, x)
     return require_insupport(μ, x), _require_insupport_pullback
 end
 
-_origin_depth_pullback(ΔΩ) = NoTangent(), NoTangent()
-ChainRulesCore.rrule(::typeof(_origin_depth), ν) = _origin_depth(ν), _origin_depth_pullback
-
 _check_dof_pullback(ΔΩ) = NoTangent(), NoTangent(), NoTangent()
 ChainRulesCore.rrule(::typeof(check_dof), ν, μ) = check_dof(ν, μ), _check_dof_pullback
 
 _checked_arg_pullback(ΔΩ) = NoTangent(), NoTangent(), ΔΩ
 ChainRulesCore.rrule(::typeof(checked_arg), ν, x) = checked_arg(ν, x), _checked_arg_pullback
 
-# = return type inference ====================================================
-
-using MeasureBase: 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
-
 end # module MeasureBaseChainRulesCoreExt

ext/MeasureBaseDistributionsChainRulesCoreExt.jl

@@ -0,0 +1,43 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseDistributionsChainRulesCoreExt
+
+using MeasureBase
+import Distributions
+import ChainRulesCore
+using ChainRulesCore: NoTangent
+
+using MeasureBase: _dist_params_numtype
+using Distributions: Distribution
+
+_dist_params_numtype_pullback(ΔΩ) = (NoTangent(), NoTangent())
+using MeasureBase: _gamma_cdf, _gamma_quantile, _beta_cdf, _beta_quantile, _gamma_logpdf, _beta_logpdf
+
+# Derivatives with respect to the variate resp. probability argument of the
+# regularized incomplete gamma and beta functions and their inverses:
+function ChainRulesCore.rrule(::typeof(_gamma_cdf), α::Real, x::Real)
+    y = _gamma_cdf(α, x)
+    dy_dx = exp(_gamma_logpdf(α, x))
+    return y, ȳ -> (NoTangent(), NoTangent(), dy_dx * ȳ)
+end
+function ChainRulesCore.rrule(::typeof(_gamma_quantile), α::Real, p::Real)
+    x = _gamma_quantile(α, p)
+    dx_dp = exp(-_gamma_logpdf(α, x))
+    return x, x̄ -> (NoTangent(), NoTangent(), dx_dp * x̄)
+end
+function ChainRulesCore.rrule(::typeof(_beta_cdf), α::Real, β::Real, x::Real)
+    y = _beta_cdf(α, β, x)
+    dy_dx = exp(_beta_logpdf(α, β, x))
+    return y, ȳ -> (NoTangent(), NoTangent(), NoTangent(), dy_dx * ȳ)
+end
+function ChainRulesCore.rrule(::typeof(_beta_quantile), α::Real, β::Real, p::Real)
+    x = _beta_quantile(α, β, p)
+    dx_dp = exp(-_beta_logpdf(α, β, x))
+    return x, x̄ -> (NoTangent(), NoTangent(), NoTangent(), dx_dp * x̄)
+end
+
+function ChainRulesCore.rrule(::typeof(_dist_params_numtype), d::Distribution)
+    _dist_params_numtype(d), _dist_params_numtype_pullback
+end
+
+end # module MeasureBaseDistributionsChainRulesCoreExt