From b699bdce5a0e75fc33a32fb4662a7710ddfc1bde Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:38:27 +0200
Subject: [PATCH 001/122] Extend static tools

---
 src/MeasureBase.jl       |   1 +
 src/combinators/power.jl |  25 +--
 src/static.jl            | 344 ++++++++++++++++++++++++++++++++++++---
 test/static.jl           | 339 +++++++++++++++++++++++++++++++++++---
 4 files changed, 648 insertions(+), 61 deletions(-)

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 2bad7d92..63149408 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -31,6 +31,7 @@ import ConstructionBase
 using ConstructionBase: constructorof
 using IntervalSets
 
+import StaticArrays
 using StaticArrays:
     StaticArray, StaticVector, StaticMatrix, SArray, SVector, SMatrix, SOneTo
 
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index e6397c3f..811e7d09 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -17,8 +17,8 @@ struct PowerMeasure{M,A} <: AbstractProductMeasure
     axes::A
 end
 
-maybestatic_length(μ::PowerMeasure) = prod(maybestatic_size(μ))
-maybestatic_size(μ::PowerMeasure) = map(maybestatic_length, μ.axes)
+maybestatic_length(μ::PowerMeasure) = size2length(maybestatic_size(μ))
+maybestatic_size(μ::PowerMeasure) = axes2size(μ.axes)
 
 function Pretty.tile(μ::PowerMeasure)
     sz = length.(μ.axes)
@@ -30,7 +30,7 @@ end
 # ToDo: Make rand return static arrays for statically-sized power measures.
 
 function _cartidxs(axs::Tuple{Vararg{AbstractUnitRange,N}}) where {N}
-    CartesianIndices(map(_dynamic, axs))
+    CartesianIndices(map(asnonstatic, axs))
 end
 
 function Base.rand(
@@ -49,11 +49,8 @@ function Base.rand(rng::AbstractRNG, ::Type{T}, d::PowerMeasure) where {T}
     end
 end
 
-@inline _pm_axes(sz::Tuple{Vararg{IntegerLike,N}}) where {N} = map(one_to, sz)
-@inline _pm_axes(axs::Tuple{Vararg{AbstractUnitRange,N}}) where {N} = axs
-
 @inline function powermeasure(x::T, sz::Tuple{Vararg{Any,N}}) where {T,N}
-    PowerMeasure(x, _pm_axes(sz))
+    PowerMeasure(x, asaxes(sz))
 end
 
 marginals(d::PowerMeasure) = fill_with(d.parent, d.axes)
@@ -86,7 +83,7 @@ for func in [:logdensityof, :logdensity_def]
         end
     end
 
-    @eval @inline function $func(d::PowerMeasure{M,Tuple{Static.SOneTo{N}}}, x) where {M,N}
+    @eval @inline function $func(d::PowerMeasure{<:Any,Tuple{<:StaticOneToLike}}, x)
         parent = d.parent
         sum(1:N) do j
             @inbounds $func(parent, x[j])
@@ -94,9 +91,9 @@ for func in [:logdensityof, :logdensity_def]
     end
 
     @eval @inline function $func(
-        d::PowerMeasure{M,NTuple{N,Static.SOneTo{0}}},
+        ::PowerMeasure{<:Any,<:Tuple{Vararg{StaticOneToLike{0}}}},
         x,
-    ) where {M,N}
+    )
         static(0.0)
     end
 end
@@ -117,11 +114,7 @@ end
     end
 end
 
-@inline getdof(μ::PowerMeasure) = getdof(μ.parent) * prod(map(length, μ.axes))
-
-@inline function getdof(::PowerMeasure{<:Any,NTuple{N,Static.SOneTo{0}}}) where {N}
-    static(0)
-end
+@inline getdof(μ::PowerMeasure) = getdof(μ.parent) * size2length(axes2size(μ.axes))
 
 @propagate_inbounds function checked_arg(μ::PowerMeasure, x::AbstractArray{<:Any})
     @boundscheck begin
@@ -144,7 +137,7 @@ logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0)
 
 # To avoid ambiguities
 function logdensity_def(
-    ::PowerMeasure{P,Tuple{Vararg{Static.SOneTo{0},N}}},
+    ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0},N}}},
     x,
 ) where {P<:PrimitiveMeasure,N}
     static(0.0)
diff --git a/src/static.jl b/src/static.jl
index da471b62..12db9585 100644
--- a/src/static.jl
+++ b/src/static.jl
@@ -1,3 +1,27 @@
+# A lots of this is about bridging Static and StaticArrays, both have their
+# own SUnitRange and SOneTo. Also provides tools to control static vs dynamic
+# array, size and axes handling.
+
+"""
+    MeasureBase.StaticUnitRange
+
+The MeasureBase default type for static unit ranges.
+"""
+const StaticUnitRange = @static if isdefined(StaticArrays, :SUnitRange)
+    # Unclear if StaticArrays.SUnitRange is part of StaticArrays stable API.
+    # Some packages use it, but let's be careful in case it disappears.
+    StaticArrays.SUnitRange
+else
+    Static.SUnitRange
+end
+
+"""
+    MeasureBase.StaticOneTo
+
+The MeasureBase default type for static one-based unit ranges.
+"""
+const StaticOneTo{T} = StaticArrays.SOneTo{T}
+
 """
     MeasureBase.IntegerLike
 
@@ -5,6 +29,67 @@ Equivalent to `Union{Integer,Static.StaticInteger}`.
 """
 const IntegerLike = Union{Integer,Static.StaticInteger}
 
+"""
+    MeasureBase.SizeLike
+
+Something that can represent the size of a collection.
+"""
+const SizeLike = Union{Tuple{},Tuple{Vararg{IntegerLike}},StaticArrays.Size}
+
+"""
+    MeasureBase.StaticSizeLike
+
+Something that can represent the size of a statically sized collection.
+"""
+const StaticSizeLike = Union{Tuple{Vararg{StaticInteger}},StaticArrays.Size}
+
+"""
+    MeasureBase.AxesLike
+
+Something that can represent axes of a collection.
+"""
+const AxesLike = Union{Tuple{},Tuple{Vararg{AbstractVector{<:IntegerLike}}}}
+
+"""
+    MeasureBase.StaticAxesLike
+
+Something that can represent axes of a statically sized collection.
+"""
+@static if isdefined(StaticArrays, :SUnitRange)
+    const StaticAxesLike = Union{
+        Tuple{Vararg{Union{StaticArrays.SOneTo,StaticArrays.SUnitRange,Static.SUnitRange}}},
+    }
+else
+    const StaticAxesLike =
+        Union{Tuple{Vararg{Union{StaticArrays.SOneTo,Static.SUnitRange}}}}
+end
+
+"""
+    const OneToLike
+
+Alias for unit ranges that start at one.
+"""
+const OneToLike = Union{Base.OneTo,StaticArrays.SOneTo,Static.SOneTo}
+
+"""
+    const StaticOneToLike{N}
+
+A static unit range from one to N.
+"""
+const StaticOneToLike{N} = Union{StaticArrays.SOneTo{N},Static.SOneTo{N}}
+
+"""
+    const StaticUnitRangeLike
+
+A static unit range.
+"""
+@static if isdefined(StaticArrays, :SUnitRange)
+    const StaticUnitRangeLike =
+        Union{StaticArrays.SOneTo,StaticArrays.SUnitRange,Static.SUnitRange}
+else
+    const StaticUnitRangeLike = Union{StaticArrays.SOneTo,Static.SUnitRange}
+end
+
 """
     MeasureBase.one_to(n::IntegerLike)
 
@@ -16,48 +101,265 @@ on the type of `n`.
 @inline one_to(n::Integer) = Base.OneTo(n)
 @inline one_to(::Static.StaticInteger{N}) where {N} = Static.SOneTo{N}()
 
-_dynamic(x::Number) = dynamic(x)
-_dynamic(::Static.SOneTo{N}) where {N} = Base.OneTo(N)
-_dynamic(r::AbstractUnitRange) = minimum(r):maximum(r)
+"""
+    MeasureBase.asnonstatic(x)
+
+Return a non-static equivalent of `x`.
+
+Defaults to `Static.dynamic(x)`.
+"""
+@inline asnonstatic(x::Number) = dynamic(x)
+@inline asnonstatic(::Tuple{}) = ()
+@static if isdefined(StaticArrays, :SUnitRange)
+    @inline asnonstatic(r::StaticArrays.SUnitRange) = r[begin]:r[end]
+end
+@inline asnonstatic(r::AbstractUnitRange) = asnonstatic(r[begin]):asnonstatic(r[end])
+@inline asnonstatic(r::Base.OneTo) = Base.OneTo(asnonstatic(r.stop))
+@inline asnonstatic(::StaticOneToLike{N}) where {N} = Base.OneTo(N)
+@inline asnonstatic(x::SizeLike) = map(asnonstatic, x)
+@inline asnonstatic(::StaticArrays.Size{TPL}) where {TPL} = TPL
+@inline asnonstatic(x::AxesLike) = map(asnonstatic, x)
 
 """
     MeasureBase.fill_with(x, sz::NTuple{N,<:IntegerLike}) where N
 
 Creates an array of size `sz` filled with `x`.
 
-Returns an instance of `FillArrays.Fill`.
+The result will typically be either a `FillArrays.Fill` or a static array,
 """
 function fill_with end
 
-@inline function fill_with(x::T, sz::Tuple{Vararg{IntegerLike,N}}) where {T,N}
-    fill_with(x, map(one_to, sz))
+@inline fill_with(x::T, n::IntegerLike) where {T} = fill_with(x, (n,))
+
+@inline fill_with(x::T, ::Tuple{}) where {T} = FillArrays.Fill(x)
+
+@inline fill_with(x, sz::SizeLike) = fill_with(x, size2axes(sz))
+
+@inline function fill_with(x::T, sz::StaticSizeLike) where {T}
+    fill(x, staticarray_type(T, canonical_size(sz)))
 end
 
-@inline function fill_with(x::T, axs::Tuple{Vararg{AbstractUnitRange,N}}) where {T,N}
-    # While `FillArrays.Fill` (mostly?) works with axes that are static unit
-    # ranges, some operations that automatic differentiation requires do fail
-    # on such instances of `Fill` (e.g. `reshape` from dynamic to static size).
-    # So need to use standard ranges for the axes for now:
-    dyn_axs = map(_dynamic, axs)
+@inline function fill_with(x, axs::AxesLike)
+    dyn_axs = map(asnonstatic, axs)
     FillArrays.Fill(x, dyn_axs)
 end
 
+# While `FillArrays.Fill` (mostly?) works with axes that are static unit
+# ranges, some operations that automatic differentiation requires do fail
+# on such instances of `Fill` (e.g. `reshape` from dynamic to static size).
+# So need to build a filled static array:
+@inline function fill_with(x::T, axs::Tuple{Vararg{StaticOneToLike}}) where {T}
+    sz = axes2size(axs)
+    fill(x, staticarray_type(T, sz))
+end
+
+"""
+    MeasureBase.staticarray_type(T, sz::StaticArrays.Size)
+
+Returns the type of a static array with element type `T` and size `sz`.
+"""
+function staticarray_type end
+
+@inline @generated function staticarray_type(
+    ::Type{T},
+    ::StaticArrays.Size{sz},
+) where {T,sz}
+    N = length(sz)
+    len = prod(sz)
+    :(SArray{Tuple{$sz...},T,$N,$len})
+end
+
+"""
+    MeasureBase.maybestatic_reshape(A, sz)
+
+Reshapes array `A` to sizes `sz`.
+
+If `A` is a static array and `sz` is static, the result is a static array.
+"""
+function maybestatic_reshape end
+
+maybestatic_reshape(A, sz) = reshape(A, canonical_size(sz))
+function maybestatic_reshape(A, sz::StaticSizeLike)
+    StaticArrays.SArray(reshape(A, canonical_size(sz)))
+end
+function maybestatic_reshape(A::StaticArray, sz::Tuple{Vararg{StaticInteger}})
+    staticarray_type(eltype(A), canonical_size(sz))(Tuple(A))
+end
+
 """
-    MeasureBase.maybestatic_length(x)::IntegerLike
+    MeasureBase.maybestatic_length(x)
 
 Returns the length of `x` as a dynamic or static integer.
 """
-maybestatic_length(x) = length(x)
-maybestatic_length(x::AbstractUnitRange) = length(x)
-function maybestatic_length(
-    ::Static.OptionallyStaticUnitRange{<:StaticInteger{A},<:StaticInteger{B}},
-) where {A,B}
-    StaticInt{B - A + 1}()
+@inline maybestatic_length(::Number) = static(1)
+@inline maybestatic_length(::Tuple{}) = static(0)
+@inline maybestatic_length(::Tuple{Vararg{Any,N}}) where {N} = static(N)
+@inline maybestatic_length(nt::NamedTuple) = maybestatic_length(values(nt))
+@inline maybestatic_length(A::AbstractArray) = size2length(maybestatic_size(A))
+@static if isdefined(StaticArrays, :SUnitRange)
+    @inline maybestatic_length(r::StaticArrays.SUnitRange) =
+        maybestatic_last(r) - maybestatic_first(r) + static(1)
 end
+@inline maybestatic_length(r::AbstractUnitRange) =
+    maybestatic_last(r) - maybestatic_first(r) + static(1)
+@inline maybestatic_length(r::Base.OneTo) = length(r)
+@inline maybestatic_length(::StaticArrays.SOneTo{N}) where {N} = static(N)
+@inline maybestatic_length(::Static.SOneTo{N}) where {N} = static(N)
 
 """
-    MeasureBase.maybestatic_size(x)::Tuple{Vararg{IntegerLike}}
+
+    MeasureBase.maybestatic_size(x)
 
 Returns the size of `x` as a tuple of dynamic or static integers.
 """
-maybestatic_size(x) = size(x)
+@inline maybestatic_size(::Number) = ()
+@inline maybestatic_size(::Tuple{}) =
+    throw(ArgumentError("Cannot determine (maybe-static) size of empty tuple"))
+@inline maybestatic_size(::Tuple{Vararg{Any,N}}) where {N} = StaticArrays.Size{(N,)}()
+@inline maybestatic_size(nt::NamedTuple) = maybestatic_size(values(nt))
+@inline maybestatic_size(A::AbstractArray) = axes2size(maybestatic_axes(A))
+@inline maybestatic_size(A::StaticArray) = StaticArrays.Size(A)
+
+"""
+    MeasureBase.maybestatic_axes(x)::Tuple{Vararg{IntegerLike}}
+
+Returns the size of `x` as a tuple of dynamic or static integers.
+"""
+@inline maybestatic_axes(::Number) = ()
+
+@inline maybestatic_axes(::Tuple{}) = (StaticOneTo(0),)
+@inline maybestatic_axes(::Tuple{Vararg{Any,N}}) where {N} = (StaticOneTo(N),)
+@inline maybestatic_axes(nt::NamedTuple) = maybestatic_axes(values(nt))
+@inline maybestatic_axes(::StaticOneToLike{N}) where {N} = (StaticOneTo(N),)
+@static if isdefined(StaticArrays, :SUnitRange)
+    @inline maybestatic_axes(r::StaticArrays.SUnitRange) = axes(r)
+end
+@inline maybestatic_axes(r::Static.OptionallyStaticUnitRange) = canonical_axes(axes(r))
+@inline maybestatic_axes(r::AbstractUnitRange) = axes(r)
+@inline maybestatic_axes(A::AbstractArray) = axes(A)
+@inline maybestatic_axes(A::StaticArray) = axes(A)
+
+"""
+    MeasureBase.axes2size(x::Tuple)
+    MeasureBase.axes2size(x::StaticArrays.Size)
+
+Get a length from a size (tuple).
+"""
+@inline axes2size(::Tuple{}) = ()
+@inline axes2size(axs::Tuple) = canonical_size(map(maybestatic_length, axs))
+
+"""map(maybestatic_length, axs)
+    MeasureBase.size2axes(sz::Tuple)
+    MeasureBase.size2axes(sz::StaticArrays.Size)
+
+Get one-based indexing axes from a size.
+"""
+@inline size2axes(::Tuple{}) = ()
+@inline size2axes(sz::Tuple) = canonical_axes(map(one_to, sz))
+@inline size2axes(::StaticArrays.Size{TPL}) where {TPL} = map(StaticOneTo, TPL)
+
+"""
+    MeasureBase.size2length(sz::Tuple)
+    MeasureBase.size2length(sz::StaticArrays.Size)
+
+Get a length from a size (tuple).
+"""
+@inline size2length(::Tuple{}) = static(1)
+@inline size2length(sz::Tuple) = prod(sz)
+@inline size2length(::StaticArrays.Size{TPL}) where {TPL} = static(prod(TPL))
+
+"""
+    MeasureBase.asaxes(axs::AxesLike)
+    MeasureBase.asaxes(sz::SizeLike)
+    MeasureBase.asaxes(len::IntegerLike)
+
+Converts axes or a size or a length of a collection to axes.
+
+One-based indexing will be used if the indexing offset can't be inferred from
+the given dimensions.
+"""
+@inline asaxes(::Tuple{}) = ()
+@inline asaxes(axs::AxesLike) = axs
+@inline asaxes(sz::SizeLike) = size2axes(sz)
+@inline asaxes(len::IntegerLike) = size2axes((len,))
+
+"""
+    MeasureBase.maybestatic_eachindex(x)
+
+Returns the the index range of `x` as a dynamic or static integer range
+"""
+maybestatic_eachindex(::Tuple{}) = StaticOneTo(0)
+maybestatic_eachindex(::Tuple{Vararg{Any,N}}) where {N} = StaticOneTo(N)
+maybestatic_eachindex(nt::NamedTuple) = maybestatic_eachindex(values(nt))
+maybestatic_eachindex(x::AbstractArray) = canonical_indices(eachindex(x))
+
+"""
+    MeasureBase.maybestatic_first(A)
+
+Returns the first element of `A` as a dynamic or static value.
+"""
+maybestatic_first(tpl::Tuple) = tpl[begin]
+maybestatic_first(nt::NamedTuple) = nt[begin]
+maybestatic_first(A::AbstractArray) = A[begin]
+maybestatic_first(::StaticArrays.Size{tpl}) where {tpl} = static(tpl[begin])
+maybestatic_first(::StaticArrays.SOneTo{N}) where {N} = static(1)
+@static if isdefined(StaticArrays, :SUnitRange)
+    maybestatic_first(::StaticArrays.SUnitRange{B,L}) where {B,L} = static(B)
+end
+function maybestatic_first(
+    ::Static.OptionallyStaticUnitRange{<:Static.StaticInteger{from},<:Static.StaticInteger},
+) where {from}
+    static(from)
+end
+
+"""
+    MeasureBase.maybestatic_last(A)
+
+Returns the last element of `A` as a dynamic or static value.
+"""
+maybestatic_last(tpl::Tuple) = tpl[end]
+maybestatic_last(nt::NamedTuple) = nt[end]
+maybestatic_last(A::AbstractArray) = A[end]
+maybestatic_last(::StaticArrays.Size{tpl}) where {tpl} = static(tpl[end])
+maybestatic_last(::StaticArrays.SOneTo{N}) where {N} = static(N)
+@static if isdefined(StaticArrays, :SUnitRange)
+    maybestatic_last(::StaticArrays.SUnitRange{B,L}) where {B,L} = static(B + L - 1)
+end
+function maybestatic_last(
+    ::Static.OptionallyStaticUnitRange{<:Any,<:Static.StaticInteger{until}},
+) where {until}
+    static(until)
+end
+
+"""
+    MeasureBase.canonical_indices(idxs::AbstractVector{<:IntegerLike})
+
+Return the canonical representation of a collection axis indices.
+"""
+@inline canonical_indices(idxs::AbstractVector{<:IntegerLike}) = idxs
+@inline canonical_indices(idxs::AbstractArray{<:CartesianIndex}) = idxs
+@inline canonical_indices(
+    ::Static.OptionallyStaticUnitRange{<:StaticInteger{1},<:StaticInteger{N}},
+) where {N} = StaticArrays.SOneTo{N}()
+@inline canonical_indices(
+    ::Static.OptionallyStaticUnitRange{<:StaticInteger{A},<:StaticInteger{B}},
+) where {A,B} = StaticUnitRange(A, B)
+@inline canonical_indices(
+    r::Static.OptionallyStaticUnitRange{<:StaticInteger{1},<:Integer},
+) = Base.OneTo(last(r))
+
+"""
+    MeasureBase.canonical_size(sz::SizeLike)
+
+Return the canonical representation of a collection size.
+"""
+@inline canonical_size(sz::SizeLike) = sz
+@inline canonical_size(sz::Tuple{Vararg{Static.StaticInteger}}) =
+    StaticArrays.Size{map(dynamic, sz)}()
+
+"""
+    MeasureBase.canonical_axes(sz::SizeLike)
+
+Return the canonical representation collection axes.
+"""
+@inline canonical_axes(axs::AxesLike) = map(canonical_indices, axs)
diff --git a/test/static.jl b/test/static.jl
index f618124b..83092ec2 100644
--- a/test/static.jl
+++ b/test/static.jl
@@ -1,34 +1,325 @@
 using Test
 
 import MeasureBase
+using MeasureBase:
+    StaticUnitRange,
+    StaticOneTo,
+    IntegerLike,
+    SizeLike,
+    StaticSizeLike,
+    AxesLike,
+    StaticAxesLike,
+    OneToLike,
+    StaticOneToLike,
+    StaticUnitRangeLike,
+    one_to,
+    asnonstatic,
+    fill_with,
+    staticarray_type,
+    maybestatic_reshape,
+    maybestatic_length,
+    maybestatic_size,
+    maybestatic_axes,
+    axes2size,
+    size2axes,
+    size2length,
+    asaxes,
+    maybestatic_eachindex,
+    maybestatic_first,
+    maybestatic_last,
+    canonical_indices,
+    canonical_size,
+    canonical_axes
 
 import Static
 using Static: static
+import StaticArrays
 import FillArrays
 
 @testset "static" begin
-    @test 2 isa MeasureBase.IntegerLike
-    @test static(2) isa MeasureBase.IntegerLike
-    @test true isa MeasureBase.IntegerLike
-    @test static(true) isa MeasureBase.IntegerLike
-
-    @test @inferred(MeasureBase.one_to(7)) isa Base.OneTo
-    @test @inferred(MeasureBase.one_to(7)) == 1:7
-    @test @inferred(MeasureBase.one_to(static(7))) isa Static.SOneTo
-    @test @inferred(MeasureBase.one_to(static(7))) == static(1):static(7)
-
-    @test @inferred(MeasureBase.fill_with(4.2, (7,))) == FillArrays.Fill(4.2, 7)
-    @test @inferred(MeasureBase.fill_with(4.2, (static(7),))) == FillArrays.Fill(4.2, 7)
-    @test @inferred(MeasureBase.fill_with(4.2, (3, static(7)))) ==
-          FillArrays.Fill(4.2, 3, 7)
-    @test @inferred(MeasureBase.fill_with(4.2, (3:7,))) == FillArrays.Fill(4.2, (3:7,))
-    @test @inferred(MeasureBase.fill_with(4.2, (static(3):static(7),))) ==
-          FillArrays.Fill(4.2, (3:7,))
-    @test @inferred(MeasureBase.fill_with(4.2, (3:7, static(2):static(5)))) ==
-          FillArrays.Fill(4.2, (3:7, 2:5))
-
-    @test MeasureBase.maybestatic_length(MeasureBase.one_to(7)) isa Int
-    @test MeasureBase.maybestatic_length(MeasureBase.one_to(7)) == 7
-    @test MeasureBase.maybestatic_length(MeasureBase.one_to(static(7))) isa Static.StaticInt
-    @test MeasureBase.maybestatic_length(MeasureBase.one_to(static(7))) == static(7)
+    v = 4.2
+    T = typeof(v)
+
+    tpl = (7, 42, 5)
+    nt = (a = 7, b = 42, c = 5)
+
+    i = 7
+    si = static(7)
+
+    @test i isa IntegerLike
+    @test si isa IntegerLike
+
+    sz = (2, 4, 3)
+    sasz = StaticArrays.Size(2, 4, 3)
+    sisz = (static(2), static(4), static(3))
+
+    len = prod(sz)
+    slen = static(len)
+
+    @test sz isa SizeLike
+    @test sasz isa SizeLike
+    @test sisz isa SizeLike
+
+    @test !(sz isa StaticSizeLike)
+    @test sasz isa StaticSizeLike
+    @test sisz isa StaticSizeLike
+
+    axs = (Base.OneTo(2), 2:5, Base.OneTo(3))
+    axs1 = (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3))
+    saaxs = (StaticOneTo(2), StaticUnitRange(2, 5), StaticOneTo(3))
+    saaxs1 = (StaticOneTo(2), StaticOneTo(4), StaticOneTo(3))
+    siaxs = (Static.SOneTo(2), static(2):static(5), static(1):static(3))
+    siaxs1 = (Static.SOneTo(2), static(1):static(4), static(1):static(3))
+
+    @test axs isa AxesLike
+    @test axs1 isa AxesLike
+    @test saaxs isa AxesLike
+    @test saaxs1 isa AxesLike
+    @test siaxs isa AxesLike
+    @test siaxs1 isa AxesLike
+
+    @test !(axs isa StaticAxesLike)
+    @test saaxs isa StaticAxesLike
+    @test saaxs1 isa StaticAxesLike
+    @test siaxs isa StaticAxesLike
+    @test siaxs1 isa StaticAxesLike
+
+    @test axs[1] isa OneToLike
+    @test !(axs[2] isa OneToLike)
+    @test axs[3] isa OneToLike
+
+    @test saaxs[1] isa OneToLike
+    @test !(saaxs[2] isa OneToLike)
+    @test saaxs1[2] isa OneToLike
+    @test saaxs[3] isa OneToLike
+
+    @test siaxs[1] isa OneToLike
+    @test !(siaxs[2] isa OneToLike)
+    @test siaxs1[2] isa OneToLike
+    @test siaxs[3] isa OneToLike
+
+    @test !(axs[1] isa StaticOneToLike)
+    @test !(axs[2] isa StaticOneToLike)
+    @test !(axs[3] isa StaticOneToLike)
+
+    @test saaxs[1] isa StaticOneToLike
+    @test !(saaxs[2] isa StaticOneToLike)
+    @test saaxs1[2] isa StaticOneToLike
+    @test saaxs[3] isa StaticOneToLike
+
+    @test siaxs[1] isa StaticOneToLike
+    @test !(siaxs[2] isa StaticOneToLike)
+    @test siaxs1[2] isa StaticOneToLike
+    @test siaxs[3] isa StaticOneToLike
+
+    @test !(axs[1] isa StaticUnitRangeLike)
+    @test !(axs[2] isa StaticUnitRangeLike)
+    @test !(axs[3] isa StaticUnitRangeLike)
+
+    @test saaxs[1] isa StaticUnitRangeLike
+    @test saaxs[2] isa StaticUnitRangeLike
+    @test saaxs1[2] isa StaticUnitRangeLike
+    @test saaxs[3] isa StaticUnitRangeLike
+
+    @test siaxs[1] isa StaticUnitRangeLike
+    @test siaxs[2] isa StaticUnitRangeLike
+    @test siaxs1[2] isa StaticUnitRangeLike
+    @test siaxs[3] isa StaticUnitRangeLike
+
+    @test @inferred(one_to(i)) == Base.OneTo(i)
+    @test @inferred(one_to(si)) == StaticOneTo(i)
+
+    @test @inferred(asnonstatic(i)) === i
+    @test @inferred(asnonstatic(si)) === i
+    @test @inferred(asnonstatic(sz)) === sz
+    @test @inferred(asnonstatic(sasz)) === sz
+    @test @inferred(asnonstatic(sisz)) === sz
+    @test @inferred(asnonstatic(axs)) === axs
+    @test @inferred(asnonstatic(saaxs)) === axs
+    @test @inferred(asnonstatic(saaxs1)) === (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3))
+    @test @inferred(asnonstatic(siaxs)) === axs
+    @test @inferred(asnonstatic(siaxs1)) === (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3))
+
+    @test @inferred(fill_with(v, i)) === FillArrays.Fill(v, i)
+    @test @inferred(fill_with(v, si)) === StaticArrays.SVector(fill(v, i)...)
+    @test @inferred(fill_with(v, ())) === FillArrays.Fill(v)
+
+    @test @inferred(fill_with(v, sz)) === FillArrays.Fill(v, sz)
+    @test @inferred(fill_with(v, sasz)) === StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz))
+    @test @inferred(fill_with(v, sisz)) === StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz))
+
+    @test @inferred(fill_with(v, axs)) === FillArrays.Fill(v, axs)
+    @test @inferred(fill_with(v, saaxs)) === FillArrays.Fill(v, axs)
+    @test @inferred(fill_with(v, saaxs1)) ===
+          StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz))
+    @test @inferred(fill_with(v, siaxs)) === FillArrays.Fill(v, axs)
+    @test @inferred(fill_with(v, siaxs1)) ===
+          StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz))
+
+    @test @inferred(staticarray_type(T, sasz)) <: StaticArrays.SArray{Tuple{2,4,3},T}
+
+    A = rand(T, len)
+    FA = FillArrays.Fill(v, len)
+    SA = StaticArrays.SVector(A...)
+
+    # Array with CartesianIndices
+    ciA = view(rand(5, 6, 6), 3:4, 2:5, 3:5)
+    ciidxs = eachindex(ciA)
+
+    rshpA = reshape(A, sz)
+    rshpFA = FillArrays.Fill(v, sz)
+    rshpSA = StaticArrays.SArray{Tuple{sz...},T}(A)
+
+    @test @inferred(maybestatic_reshape(A, sz)) == rshpA
+    @test typeof(maybestatic_reshape(A, sz)) == typeof(rshpA)
+    @test @inferred(maybestatic_reshape(A, sasz)) == rshpA
+    @test maybestatic_reshape(A, sasz) isa StaticArrays.SArray
+    @test @inferred(maybestatic_reshape(A, sisz)) == rshpA
+    @test maybestatic_reshape(A, sisz) isa StaticArrays.SArray
+
+    @test @inferred(maybestatic_reshape(FA, sz)) == rshpFA
+    @test typeof(maybestatic_reshape(FA, sz)) == typeof(rshpFA)
+    @test @inferred(maybestatic_reshape(FA, sasz)) == rshpFA
+    @test maybestatic_reshape(FA, sasz) isa StaticArrays.SArray
+    @test @inferred(maybestatic_reshape(FA, sisz)) == rshpFA
+    @test maybestatic_reshape(FA, sisz) isa StaticArrays.SArray
+
+    @test @inferred(maybestatic_reshape(SA, sz)) == rshpA
+    @test maybestatic_reshape(SA, sz) isa Base.ReshapedArray{T,3,<:StaticArrays.SVector}
+    @test @inferred(maybestatic_reshape(SA, sasz)) === rshpSA
+    @test @inferred(maybestatic_reshape(SA, sisz)) === rshpSA
+
+    @test @inferred(maybestatic_length(5)) === static(1)
+    @test @inferred(maybestatic_length(())) === static(0)
+    @test @inferred(maybestatic_length((sz))) === static(3)
+    @test @inferred(maybestatic_length((a = 2, b = 4, c = 3))) === static(3)
+    @test @inferred(maybestatic_length(Base.OneTo(4))) === 4
+    @test @inferred(maybestatic_length(StaticArrays.SOneTo(4))) === static(4)
+    @test @inferred(maybestatic_length(Static.SOneTo(4))) === static(4)
+    @test @inferred(maybestatic_length(static(2):static(5))) === static(4)
+    @test @inferred(maybestatic_length(rshpA)) === length(rshpA)
+    @test @inferred(maybestatic_length(rshpFA)) === length(rshpA)
+    @test @inferred(maybestatic_length(rshpSA)) === static(length(rshpA))
+
+    @test @inferred(maybestatic_size(5)) === ()
+    @test_throws ArgumentError maybestatic_size(())
+    @test @inferred(maybestatic_size((sz))) === StaticArrays.Size(3)
+    @test @inferred(maybestatic_size((a = 2, b = 4, c = 3))) === StaticArrays.Size(3)
+    @test @inferred(maybestatic_size(Base.OneTo(4))) === (4,)
+    @test @inferred(maybestatic_size(StaticArrays.SOneTo(4))) === StaticArrays.Size(4)
+    @test @inferred(maybestatic_size(StaticUnitRange(2, 5))) === StaticArrays.Size(4)
+    @test @inferred(maybestatic_size(Static.SOneTo(4))) === StaticArrays.Size(4)
+    @test @inferred(maybestatic_size(static(2):static(5))) === StaticArrays.Size(4)
+    @test @inferred(maybestatic_size(rshpA)) === size(rshpA)
+    @test @inferred(maybestatic_size(rshpFA)) === size(rshpA)
+    @test @inferred(maybestatic_size(rshpSA)) === StaticArrays.Size(size(rshpA)...)
+
+    @test @inferred(maybestatic_axes(5)) === ()
+    @test @inferred(maybestatic_axes(())) === (StaticOneTo(0),)
+    @test @inferred(maybestatic_axes((sz))) === (StaticOneTo(3),)
+    @test @inferred(maybestatic_axes((a = 2, b = 4, c = 3))) === (StaticOneTo(3),)
+    @test @inferred(maybestatic_axes(Base.OneTo(4))) === (Base.OneTo(4),)
+    @test @inferred(maybestatic_axes(StaticArrays.SOneTo(4))) === (StaticOneTo(4),)
+    @test @inferred(maybestatic_axes(Static.SOneTo(4))) === (StaticOneTo(4),)
+    @test @inferred(maybestatic_axes(static(2):static(5))) === (StaticOneTo(4),)
+    @test @inferred(maybestatic_axes(rshpA)) === axes(rshpA)
+    @test @inferred(maybestatic_axes(rshpFA)) === axes(rshpA)
+    @test @inferred(maybestatic_axes(rshpSA)) === saaxs1
+
+    @test @inferred(axes2size(())) === ()
+    @test @inferred(axes2size(axs)) === sz
+    @test @inferred(axes2size(saaxs)) === sasz
+    @test @inferred(axes2size(saaxs1)) === sasz
+    @test @inferred(axes2size(siaxs)) === sasz
+    @test @inferred(axes2size(siaxs1)) === sasz
+
+    @test @inferred(size2axes(())) === ()
+    @test @inferred(size2axes(sz)) === axs1
+    @test @inferred(size2axes(sasz)) === saaxs1
+    @test @inferred(size2axes(sisz)) === saaxs1
+
+    @test @inferred(size2length(())) === static(1)
+    @test @inferred(size2length(sz)) === len
+    @test @inferred(size2length(sasz)) === slen
+    @test @inferred(size2length(sisz)) === slen
+
+    @test @inferred(asaxes(())) === ()
+    @test @inferred(asaxes(len)) === (Base.OneTo(len),)
+    @test @inferred(asaxes(slen)) === (StaticOneTo(len),)
+    @test @inferred(asaxes(sz)) === axs1
+    @test @inferred(asaxes(sasz)) === saaxs1
+    @test @inferred(asaxes(sisz)) === saaxs1
+    @test @inferred(asaxes(axs)) === axs
+    @test @inferred(asaxes(axs1)) === axs1
+    @test @inferred(asaxes(saaxs)) === saaxs
+    @test @inferred(asaxes(saaxs1)) === saaxs1
+    @test @inferred(asaxes(siaxs)) === siaxs
+    @test @inferred(asaxes(siaxs1)) === siaxs1
+
+    @test @inferred(maybestatic_eachindex(())) === StaticOneTo(0)
+    @test @inferred(maybestatic_eachindex(tpl)) === StaticOneTo(3)
+    @test @inferred(maybestatic_eachindex(nt)) === StaticOneTo(3)
+    @test @inferred(maybestatic_eachindex(axs[1])) === Base.OneTo(length(axs[1]))
+    @test @inferred(maybestatic_eachindex(axs[2])) === Base.OneTo(length(axs[2]))
+    @test @inferred(maybestatic_eachindex(saaxs[1])) === StaticOneTo(length(axs[1]))
+    @test @inferred(maybestatic_eachindex(saaxs[2])) === StaticOneTo(length(axs[2]))
+    @test @inferred(maybestatic_eachindex(siaxs[1])) === StaticOneTo(length(axs[1]))
+    @test @inferred(maybestatic_eachindex(siaxs[2])) === StaticOneTo(length(axs[2]))
+    @test @inferred(maybestatic_eachindex(A)) === Base.OneTo(24)
+    @test @inferred(maybestatic_eachindex(ciA)) === eachindex(ciA)
+    @test @inferred(maybestatic_eachindex(FA)) === Base.OneTo(24)
+    @test @inferred(maybestatic_eachindex(SA)) === StaticOneTo(24)
+
+    @test_throws BoundsError maybestatic_first(())
+    @test @inferred(maybestatic_first(tpl)) === first(tpl)
+    @test @inferred(maybestatic_first(nt)) === first(nt)
+    @test @inferred(maybestatic_first(sz)) === first(sz)
+    @test @inferred(maybestatic_first(sasz)) === static(first(sz))
+    @test @inferred(maybestatic_first(sisz)) === static(first(sz))
+    @test @inferred(maybestatic_first(axs[1])) === first(axs[1])
+    @test @inferred(maybestatic_first(axs[2])) === first(axs[2])
+    @test @inferred(maybestatic_first(saaxs[1])) === static(first(axs[1]))
+    @test @inferred(maybestatic_first(saaxs[2])) === static(first(axs[2]))
+    @test @inferred(maybestatic_first(siaxs[1])) === static(first(axs[1]))
+    @test @inferred(maybestatic_first(siaxs[2])) === static(first(axs[2]))
+    @test @inferred(maybestatic_first(A)) === first(A)
+    @test @inferred(maybestatic_first(ciA)) === first(ciA)
+    @test @inferred(maybestatic_first(FA)) === first(FA)
+    @test @inferred(maybestatic_first(SA)) === first(SA)
+
+    @test_throws BoundsError maybestatic_last(())
+    @test @inferred(maybestatic_last(tpl)) === last(tpl)
+    @test @inferred(maybestatic_last(nt)) === last(nt)
+    @test @inferred(maybestatic_last(sz)) === last(sz)
+    @test @inferred(maybestatic_last(sasz)) === static(last(sz))
+    @test @inferred(maybestatic_last(sisz)) === static(last(sz))
+    @test @inferred(maybestatic_last(axs[1])) === last(axs[1])
+    @test @inferred(maybestatic_last(axs[2])) === last(axs[2])
+    @test @inferred(maybestatic_last(saaxs[1])) === static(last(axs[1]))
+    @test @inferred(maybestatic_last(saaxs[2])) === static(last(axs[2]))
+    @test @inferred(maybestatic_last(siaxs[1])) === static(last(axs[1]))
+    @test @inferred(maybestatic_last(siaxs[2])) === static(last(axs[2]))
+    @test @inferred(maybestatic_last(A)) === last(A)
+    @test @inferred(maybestatic_last(ciA)) === last(ciA)
+    @test @inferred(maybestatic_last(FA)) === last(FA)
+    @test @inferred(maybestatic_last(SA)) === last(SA)
+
+    @test @inferred(canonical_indices(axs[1])) === axs[1]
+    @test @inferred(canonical_indices(axs[2])) === axs[2]
+    @test @inferred(canonical_indices(saaxs[1])) === saaxs[1]
+    @test @inferred(canonical_indices(saaxs[2])) === saaxs[2]
+    @test @inferred(canonical_indices(siaxs[1])) === saaxs[1]
+    @test @inferred(canonical_indices(siaxs[2])) === saaxs[2]
+    @test @inferred(canonical_indices(ciidxs)) === ciidxs
+
+    @test @inferred(canonical_size(sz)) === sz
+    @test @inferred(canonical_size(sasz)) === sasz
+    @test @inferred(canonical_size(sisz)) === sasz
+
+    @test @inferred(canonical_axes(axs)) === axs
+    @test @inferred(canonical_axes(axs1)) === axs1
+    @test @inferred(canonical_axes(saaxs)) === saaxs
+    @test @inferred(canonical_axes(saaxs1)) === saaxs1
+    @test @inferred(canonical_axes(siaxs)) === saaxs
+    @test @inferred(canonical_axes(siaxs1)) === saaxs1
 end

From 1eca09e9dcf43dd3cb3a1f630bc441cbd4c968e0 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:38:27 +0200
Subject: [PATCH 002/122] Add internal infer_logdensity_type

Shouldn't call Core.Compiler.return_type directly in many places.
---
 src/utils.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/utils.jl b/src/utils.jl
index 0ec81a50..5d05d8b1 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -133,6 +133,11 @@ function infer_zero(f, args...)
     zero(typeintersect(AbstractFloat, inferred_type))
 end
 
+function infer_logdensity_type(f::F, ::M, ::Type{T}) where {F,M,T}
+    inferred_type = Core.Compiler.return_type(f, Tuple{M,T})
+    return inferred_type
+end
+
 @inline function allequal(f, x::AbstractArray)
     val = f(first(x))
     @simd for xj in x

From 970ebdbb36635eab4bdb964835075fcb942abcb5 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:38:27 +0200
Subject: [PATCH 003/122] Make logdensityof for PowerMeasure handle empty
 powers/variates

---
 src/combinators/power.jl | 15 ++++++++++++---
 test/test_basics.jl      |  5 +++--
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 811e7d09..5c0b031d 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -77,9 +77,18 @@ end
 
 for func in [:logdensityof, :logdensity_def]
     @eval @inline function $func(d::PowerMeasure{M}, x) where {M}
-        parent = d.parent
-        sum(x) do xj
-            $func(parent, xj)
+        parent_m = d.parent
+        sz_parent = axes2size(d.axes)
+        sz_x = maybestatic_size(x)
+        if sz_parent != sz_x
+            throw(ArgumentError("Size of variate doesn't match size of power measure"))
+        end
+        R = infer_logdensity_type($func, parent_m, eltype(x))
+        if isempty(x)
+            return zero(R)::R
+        else
+            # Need to convert since sum can turn static into dynamic values:
+            return convert(R, sum(Base.Fix1($func, parent_m), x))::R
         end
     end
 
diff --git a/test/test_basics.jl b/test/test_basics.jl
index 7ac29dc1..bd5a409c 100644
--- a/test/test_basics.jl
+++ b/test/test_basics.jl
@@ -120,8 +120,9 @@ end
 end
 
 @testset "powers" begin
-    @test logdensityof(Lebesgue()^3, 2) == logdensityof(Lebesgue()^(3,), 2)
-    @test logdensityof(Lebesgue()^3, 2) == logdensityof(Lebesgue()^(3, 1), (2, 0))
+    @test logdensityof(Lebesgue()^3, [2, 2, 2]) == logdensityof(Lebesgue()^(3,), fill(2, 3))
+    @test logdensityof(Lebesgue()^3, fill(2, 3)) ==
+          logdensityof(Lebesgue()^(3, 1), fill(2, 3, 1))
 end
 
 NormalMeasure() = ∫exp(x -> -0.5x^2, Lebesgue(ℝ))

From ce8d501121ead86f4bc426fcc8b0dea8375aa13f Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:38:27 +0200
Subject: [PATCH 004/122] Use internal  _TransportToStd as a function

---
 src/standard/stdmeasure.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index 833f280e..4b957651 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -57,7 +57,7 @@ end
 # Helpers for product transforms and similar:
 
 struct _TransportToStd{NU<:StdMeasure} <: Function end
-_TransportToStd{NU}(μ, x) where {NU} = transport_to(NU()^getdof(μ), μ)(x)
+(::_TransportToStd{NU})(μ, x) where {NU} = transport_to(NU()^getdof(μ), μ)(x)
 
 struct _TransportFromStd{MU<:StdMeasure} <: Function end
 _TransportFromStd{MU}(ν, x) where {MU} = transport_to(ν, MU()^getdof(ν))(x)
@@ -67,7 +67,7 @@ function _tuple_transport_def(
     μs::Tuple,
     xs::Tuple,
 ) where {NU<:StdMeasure}
-    reshape(vcat(map(_TransportToStd{NU}, μs, xs)...), ν.axes)
+    reshape(vcat(map(_TransportToStd{NU}(), μs, xs)...), ν.axes)
 end
 
 function transport_def(
@@ -93,7 +93,7 @@ end
 function _stdvar_viewranges(μs::Tuple, startidx::IntegerLike)
     N = map(getdof, μs)
     offs = _offset_cumsum(startidx, N...)
-    map((o, n) -> o:o+n-1, offs, N)
+    map((o, n) -> o:(o+n-1), offs, N)
 end
 
 function _tuple_transport_def(

From 52a610ba1321fe5bab71e47f419180c075160c82 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:38:27 +0200
Subject: [PATCH 005/122] Add pwr_base, pwr_axes, pwr_size

---
 src/combinators/power.jl | 35 ++++++++++++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 5 deletions(-)

diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 5c0b031d..62b4336f 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -11,6 +11,8 @@ the product determines the dimensionality of the resulting support.
 Note that power measures are only well-defined for integer powers.
 
 The nth power of a measure μ can be written μ^n.
+
+See also [`pwr_base`](@ref), [`pwr_axes`](@ref) and [`pwr_size`](@ref).
 """
 struct PowerMeasure{M,A} <: AbstractProductMeasure
     parent::M
@@ -20,6 +22,27 @@ end
 maybestatic_length(μ::PowerMeasure) = size2length(maybestatic_size(μ))
 maybestatic_size(μ::PowerMeasure) = axes2size(μ.axes)
 
+"""
+    MeasureBase.pwr_base(μ::PowerMeasure)
+
+Returns `ν` for `μ = ν^axs`
+"""
+@inline pwr_base(μ::PowerMeasure) = μ.parent
+
+"""
+    MeasureBase.pwr_axes(μ::PowerMeasure)
+
+Returns `axs` for `μ = ν^axs`, `axs` being a tuple of integer ranges.
+"""
+@inline pwr_axes(μ::PowerMeasure) = μ.axes
+
+"""
+    MeasureBase.pwr_size(μ::PowerMeasure)
+
+Returns `sz` for `μ = ν^sz`, `sz` being a tuple of integers.
+"""
+@inline pwr_size(μ::PowerMeasure) = axes2size(μ.axes)
+
 function Pretty.tile(μ::PowerMeasure)
     sz = length.(μ.axes)
     arg1 = Pretty.tile(μ.parent)
@@ -38,14 +61,16 @@ function Base.rand(
     ::Type{T},
     d::PowerMeasure{M},
 ) where {T,M<:AbstractMeasure}
-    map(_cartidxs(d.axes)) do _
-        rand(rng, T, d.parent)
+    axs, base_d = pwr_axes(d), pwr_base(d)
+    map(_cartidxs(axs)) do _
+        rand(rng, T, base_d)
     end
 end
 
 function Base.rand(rng::AbstractRNG, ::Type{T}, d::PowerMeasure) where {T}
-    map(_cartidxs(d.axes)) do _
-        rand(rng, d.parent)
+    axs, base_d = pwr_axes(d), pwr_base(d)
+    map(_cartidxs(axs)) do _
+        rand(rng, base_d)
     end
 end
 
@@ -127,7 +152,7 @@ end
 
 @propagate_inbounds function checked_arg(μ::PowerMeasure, x::AbstractArray{<:Any})
     @boundscheck begin
-        sz_μ = map(length, μ.axes)
+        sz_μ = pwr_size(μ)
         sz_x = size(x)
         if sz_μ != sz_x
             throw(ArgumentError("Size of variate doesn't match size of power measure"))

From d4d3100620648ad9b0bc76a9bdd0bbe02c5ebd36 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:38:27 +0200
Subject: [PATCH 006/122] Code formatting

---
 src/combinators/implicitlymapped.jl | 4 ++--
 src/density-core.jl                 | 4 ++--
 src/interface.jl                    | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/combinators/implicitlymapped.jl b/src/combinators/implicitlymapped.jl
index 964ea466..3966b10a 100644
--- a/src/combinators/implicitlymapped.jl
+++ b/src/combinators/implicitlymapped.jl
@@ -179,13 +179,13 @@ struct TakeAny{T<:IntegerLike}
     n::T
 end
 
-_takeany_range(f::TakeAny, idxs) = first(idxs):first(idxs)+dynamic(f.n)-1
+_takeany_range(f::TakeAny, idxs) = first(idxs):(first(idxs)+dynamic(f.n)-1)
 @inline _takeany_range(f::TakeAny, ::OneTo) = OneTo(dynamic(f.n))
 
 @inline _takeany_range(::TakeAny{<:Static.StaticInteger{N}}, ::OneTo) where {N} = SOneTo(N)
 @inline _takeany_range(::TakeAny{<:Static.StaticInteger{N}}, ::SOneTo) where {N} = SOneTo(N)
 
-@inline (f::TakeAny)(xs::Tuple) = xs[begin:begin+f.n-1]
+@inline (f::TakeAny)(xs::Tuple) = xs[begin:(begin+f.n-1)]
 @inline (f::TakeAny)(xs::AbstractVector) = xs[_takeany_range(f, eachindex(xs))]
 
 function (f::TakeAny)(xs)
diff --git a/src/density-core.jl b/src/density-core.jl
index 6ac3d01e..f3b2db2b 100644
--- a/src/density-core.jl
+++ b/src/density-core.jl
@@ -149,13 +149,13 @@ end
         ℓ = logdensity_def(μs[$M], νs[$N], x)
     end
 
-    for i in 1:M-1
+    for i in 1:(M-1)
         push!(q.args, :(Δℓ = logdensity_def(μs[$i], x)))
         # push!(q.args, :(println("Adding", Δℓ)))
         push!(q.args, :(ℓ += Δℓ))
     end
 
-    for j in 1:N-1
+    for j in 1:(N-1)
         push!(q.args, :(Δℓ = logdensity_def(νs[$j], x)))
         # push!(q.args, :(println("Subtracting", Δℓ)))
         push!(q.args, :(ℓ -= Δℓ))
diff --git a/src/interface.jl b/src/interface.jl
index 18080ac7..4890ddd6 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -110,7 +110,7 @@ function test_smf(μ, n = 100)
     @testset "smf($μ)" begin
         # Get `n` sorted uniforms in O(n) time
         p = rand(n)
-        p .+= 0:n-1
+        p .+= 0:(n-1)
         p .*= inv(n)
 
         F(x) = smf(μ, x)

From 6f6087d620a1a3df57700f71224162c4aca0b0c9 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:38:27 +0200
Subject: [PATCH 007/122] Add ForwardDiff extension

---
 Project.toml                     |  3 +++
 ext/MeasureBaseForwardDiffExt.jl | 14 ++++++++++++++
 2 files changed, 17 insertions(+)
 create mode 100644 ext/MeasureBaseForwardDiffExt.jl

diff --git a/Project.toml b/Project.toml
index 44f89f80..501028be 100644
--- a/Project.toml
+++ b/Project.toml
@@ -33,9 +33,11 @@ Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 
 [extensions]
 MeasureBaseChainRulesCoreExt = "ChainRulesCore"
+MeasureBaseForwardDiffExt = "ForwardDiff"
 
 [compat]
 ChainRulesCore = "1"
@@ -45,6 +47,7 @@ ConstantRNGs = "0.1.1"
 ConstructionBase = "1.3"
 DensityInterface = "0.4"
 FillArrays = "0.12, 0.13, 1"
+ForwardDiff = "0.8, 0.9, 0.10"
 FunctionChains = "0.2"
 IfElse = "0.1"
 IntervalSets = "0.7"
diff --git a/ext/MeasureBaseForwardDiffExt.jl b/ext/MeasureBaseForwardDiffExt.jl
new file mode 100644
index 00000000..8a1cab44
--- /dev/null
+++ b/ext/MeasureBaseForwardDiffExt.jl
@@ -0,0 +1,14 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseForwardDiffExt
+
+using MeasureBase
+import ForwardDiff
+
+function MeasureBase.containsnan(x::ForwardDiff.Dual)
+    a = containsnan(x.value)
+    b = containsnan(x.partials)
+    return a || b
+end
+
+end # module MeasureBaseForwardDiffExt

From 4986a9ae5266fd844f6df4d5ce4453d2e37113cb Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:04 +0200
Subject: [PATCH 008/122] Add Distributions and DistributionsForwardDiff
 extensions

---
 Project.toml                                  | 5 +++++
 ext/MeasureBaseDistributionsExt.jl            | 8 ++++++++
 ext/MeasureBaseDistributionsForwardDiffExt.jl | 9 +++++++++
 3 files changed, 22 insertions(+)
 create mode 100644 ext/MeasureBaseDistributionsExt.jl
 create mode 100644 ext/MeasureBaseDistributionsForwardDiffExt.jl

diff --git a/Project.toml b/Project.toml
index 501028be..17f621dc 100644
--- a/Project.toml
+++ b/Project.toml
@@ -33,10 +33,13 @@ Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 
 [extensions]
 MeasureBaseChainRulesCoreExt = "ChainRulesCore"
+MeasureBaseDistributionsExt = "Distributions"
+MeasureBaseDistributionsForwardDiffExt = ["Distributions", "ForwardDiff"]
 MeasureBaseForwardDiffExt = "ForwardDiff"
 
 [compat]
@@ -46,6 +49,8 @@ Compat = "3.35, 4"
 ConstantRNGs = "0.1.1"
 ConstructionBase = "1.3"
 DensityInterface = "0.4"
+Distributions = "0.25.1"
+Distributions = "0.25.111"
 FillArrays = "0.12, 0.13, 1"
 ForwardDiff = "0.8, 0.9, 0.10"
 FunctionChains = "0.2"
diff --git a/ext/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt.jl
new file mode 100644
index 00000000..beb47821
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt.jl
@@ -0,0 +1,8 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseDistributionsExt
+
+using MeasureBase
+import Distributions
+
+end # module MeasureBaseDistributionsExt
diff --git a/ext/MeasureBaseDistributionsForwardDiffExt.jl b/ext/MeasureBaseDistributionsForwardDiffExt.jl
new file mode 100644
index 00000000..36218eec
--- /dev/null
+++ b/ext/MeasureBaseDistributionsForwardDiffExt.jl
@@ -0,0 +1,9 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseDistributionsForwardDiffExt
+
+using MeasureBase
+import Distributions
+import ForwardDiff
+
+end # module MeasureBaseDistributionsForwardDiffExt

From a866672c0a4503b6adbd6281577f10c9fea473ac Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:10 +0200
Subject: [PATCH 009/122] Add DistributionsChainRulesCore extension

---
 Project.toml                                     | 1 +
 ext/MeasureBaseDistributionsChainRulesCoreExt.jl | 9 +++++++++
 2 files changed, 10 insertions(+)
 create mode 100644 ext/MeasureBaseDistributionsChainRulesCoreExt.jl

diff --git a/Project.toml b/Project.toml
index 17f621dc..06a51ed9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -39,6 +39,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 [extensions]
 MeasureBaseChainRulesCoreExt = "ChainRulesCore"
 MeasureBaseDistributionsExt = "Distributions"
+MeasureBaseDistributionsChainRulesCoreExt = ["Distributions", "ChainRulesCore"]
 MeasureBaseDistributionsForwardDiffExt = ["Distributions", "ForwardDiff"]
 MeasureBaseForwardDiffExt = "ForwardDiff"
 
diff --git a/ext/MeasureBaseDistributionsChainRulesCoreExt.jl b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl
new file mode 100644
index 00000000..4dd3f4ff
--- /dev/null
+++ b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl
@@ -0,0 +1,9 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseDistributionsChainRulesCoreExt
+
+using MeasureBase
+import Distributions
+import ChainRulesCore
+
+end # module MeasureBaseDistributionsChainRulesCoreExt

From 88fe1680a62df947c7574d10f4b003048e56b001 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:10 +0200
Subject: [PATCH 010/122] Add function asmeasure

Will be used a lot when bridging from Distributions to MeasureBase.
---
 src/MeasureBase.jl | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 63149408..b871fd43 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -65,6 +65,21 @@ abstract type AbstractMeasure end
 
 AbstractMeasure(m::AbstractMeasure) = m
 
+
+"""
+    asmeasure(m)
+
+Turns a measure-like object `m` into an `AbstractMeasure`.
+
+Calls `convert(AbstractMeasure, m)` by default
+"""
+function asmeasure end
+
+@inline asmeasure(m::AbstractMeasure) = m
+asmeasure(m) = convert(AbstractMeasure, m)
+export asmeasure
+
+
 function Pretty.quoteof(d::M) where {M<:AbstractMeasure}
     the_names = fieldnames(typeof(d))
     :($M($([getfield(d, n) for n in the_names]...)))

From e8fd530a7ee7fcb84189420775d4405b8a715b63 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:10 +0200
Subject: [PATCH 011/122] Add AsMeasure

---
 src/MeasureBase.jl | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index b871fd43..26eab64f 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -65,7 +65,6 @@ abstract type AbstractMeasure end
 
 AbstractMeasure(m::AbstractMeasure) = m
 
-
 """
     asmeasure(m)
 
@@ -79,6 +78,25 @@ function asmeasure end
 asmeasure(m) = convert(AbstractMeasure, m)
 export asmeasure
 
+"""
+    struct AsMeasure{T}
+
+Wrapes a measure-like object into an `AbstractMeasure`.
+
+Constructor:
+
+```
+AsMeasure{T}(obj::T)
+```
+
+User code should not create instances of `AsMeasure` directly, but should
+call `asmeasure(obj)` instead.
+"""
+struct AsMeasure{T} <: AbstractMeasure
+    obj::T
+
+    AsMeasure{T}(obj::T) where {T} = new(obj)
+end
 
 function Pretty.quoteof(d::M) where {M<:AbstractMeasure}
     the_names = fieldnames(typeof(d))

From 73d29a1595848b3f76469dbdf30df933bf1f00f3 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:10 +0200
Subject: [PATCH 012/122] Add collection utils

---
 ext/MeasureBaseChainRulesCoreExt.jl | 44 +++++++++++++++++++++++++++++
 src/MeasureBase.jl                  |  1 +
 src/collection_utils.jl             | 24 ++++++++++++++++
 3 files changed, 69 insertions(+)
 create mode 100644 src/collection_utils.jl

diff --git a/ext/MeasureBaseChainRulesCoreExt.jl b/ext/MeasureBaseChainRulesCoreExt.jl
index 57ed25fa..0384a04b 100644
--- a/ext/MeasureBaseChainRulesCoreExt.jl
+++ b/ext/MeasureBaseChainRulesCoreExt.jl
@@ -16,6 +16,50 @@ ChainRulesCore.rrule(::typeof(isneginf), x) = isneginf(x), _logdensityof_rt_pull
 _isposinf_pullback(::Any) = (NoTangent(), ZeroTangent())
 ChainRulesCore.rrule(::typeof(isposinf), x) = isposinf(x), _isposinf_pullback
 
+# = collection utils =========================================================
+
+using MeasureBase: _dropfront, _dropback, _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
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 26eab64f..7957bb1a 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -148,6 +148,7 @@ using IrrationalConstants
 using IrrationalConstants: loghalf
 
 include("static.jl")
+include("collection_utils.jl")
 include("smf.jl")
 include("getdof.jl")
 include("transport.jl")
diff --git a/src/collection_utils.jl b/src/collection_utils.jl
new file mode 100644
index 00000000..1de51f7e
--- /dev/null
+++ b/src/collection_utils.jl
@@ -0,0 +1,24 @@
+function _pushfront(v::AbstractVector, x)
+    T = promote_type(eltype(v), typeof(x))
+    r = similar(v, T, length(eachindex(v)) + 1)
+    r[firstindex(r)] = x
+    r[firstindex(r)+1:lastindex(r)] = v
+    r
+end
+
+function _pushback(v::AbstractVector, x)
+    T = promote_type(eltype(v), typeof(x))
+    r = similar(v, T, length(eachindex(v)) + 1)
+    r[lastindex(r)] = x
+    r[firstindex(r):lastindex(r)-1] = v
+    r
+end
+
+_dropfront(v::AbstractVector) = v[firstindex(v)+1:lastindex(v)]
+
+_dropback(v::AbstractVector) = v[firstindex(v):lastindex(v)-1]
+
+_rev_cumsum(xs::AbstractVector) = reverse(cumsum(reverse(xs)))
+
+# Equivalent to `cumprod(xs)``:
+_exp_cumsum_log(xs::AbstractVector) = exp.(cumsum(log.(xs)))

From 4e0645d23e707b2a90aeb3f029f7c48b277182f9 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:10 +0200
Subject: [PATCH 013/122] Require FunctionChains v0.2.3

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 06a51ed9..5a31c173 100644
--- a/Project.toml
+++ b/Project.toml
@@ -54,7 +54,7 @@ Distributions = "0.25.1"
 Distributions = "0.25.111"
 FillArrays = "0.12, 0.13, 1"
 ForwardDiff = "0.8, 0.9, 0.10"
-FunctionChains = "0.2"
+FunctionChains = "0.2.3"
 IfElse = "0.1"
 IntervalSets = "0.7"
 InverseFunctions = "0.1.8"

From 19dd55d77d2c3628902a69025d03c8a4c18c63d5 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:10 +0200
Subject: [PATCH 014/122] Require HeterogeneousComputing

---
 Project.toml       | 2 ++
 src/MeasureBase.jl | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/Project.toml b/Project.toml
index 5a31c173..fd3f862a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -11,6 +11,7 @@ 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"
@@ -55,6 +56,7 @@ Distributions = "0.25.111"
 FillArrays = "0.12, 0.13, 1"
 ForwardDiff = "0.8, 0.9, 0.10"
 FunctionChains = "0.2.3"
+HeterogeneousComputing = "0.2.3"
 IfElse = "0.1"
 IntervalSets = "0.7"
 InverseFunctions = "0.1.8"
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 7957bb1a..e1a3dd12 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -44,6 +44,9 @@ using Static: StaticInteger
 using FunctionChains
 using PropertyFunctions: PropSelFunction
 
+import HeterogeneousComputing
+using HeterogeneousComputing: real_numtype
+
 export gentype
 export rebase
 

From 1d52fdd587809d0e0ef144bbb235600abbb0133b Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:10 +0200
Subject: [PATCH 015/122] Use StaticThings.jl for static-size tooling

---
 Project.toml               |   2 +
 src/MeasureBase.jl         |  14 +-
 src/combinators/power.jl   |  17 +-
 src/standard/stdmeasure.jl |   2 +-
 src/static.jl              | 365 -------------------------------------
 test/static.jl             | 325 ---------------------------------
 6 files changed, 26 insertions(+), 699 deletions(-)
 delete mode 100644 src/static.jl
 delete mode 100644 test/static.jl

diff --git a/Project.toml b/Project.toml
index fd3f862a..7dea2fc6 100644
--- a/Project.toml
+++ b/Project.toml
@@ -28,6 +28,7 @@ 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"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
@@ -73,6 +74,7 @@ Reexport = "1"
 SpecialFunctions = "2"
 Static = "0.8, 1"
 StaticArrays = "1.5"
+StaticThings = "0.2"
 Statistics = "1"
 Test = "1"
 Tricks = "0.1"
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index e1a3dd12..c4c0d955 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -44,6 +44,19 @@ using Static: StaticInteger
 using FunctionChains
 using PropertyFunctions: PropSelFunction
 
+using StaticThings:
+    AxesLike, StaticAxesLike, SizeLike, StaticSizeLike,
+    OneToLike, StaticOneTo, StaticOneToLike, RealLike,
+    IntegerLike, StaticUnitRange, StaticUnitRangeLike,
+    NoTypeSize,
+    asaxes, asnonstatic, 
+    canonical_axes, canonical_indices, canonical_size,
+    maybestatic_axes, maybestatic_eachindex,
+    maybestatic_length, maybestatic_size, maybestatic_first, maybestatic_last,
+    maybestatic_oneto, maybestatic_fill, maybestatic_reshape,
+    size_from_type, axes2size, size2axes, size2length,
+    staticarray_type
+
 import HeterogeneousComputing
 using HeterogeneousComputing: real_numtype
 
@@ -150,7 +163,6 @@ using Compat
 using IrrationalConstants
 using IrrationalConstants: loghalf
 
-include("static.jl")
 include("collection_utils.jl")
 include("smf.jl")
 include("getdof.jl")
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 62b4336f..4d4760ba 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -1,4 +1,5 @@
 import Base
+import StaticThings
 
 export PowerMeasure
 
@@ -19,8 +20,8 @@ struct PowerMeasure{M,A} <: AbstractProductMeasure
     axes::A
 end
 
-maybestatic_length(μ::PowerMeasure) = size2length(maybestatic_size(μ))
-maybestatic_size(μ::PowerMeasure) = axes2size(μ.axes)
+StaticThings.maybestatic_length(μ::PowerMeasure) = size2length(maybestatic_size(μ))
+StaticThings.maybestatic_size(μ::PowerMeasure) = axes2size(μ.axes)
 
 """
     MeasureBase.pwr_base(μ::PowerMeasure)
@@ -78,13 +79,13 @@ end
     PowerMeasure(x, asaxes(sz))
 end
 
-marginals(d::PowerMeasure) = fill_with(d.parent, d.axes)
+marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes)
 
 function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{AbstractArray,N}}) where {N}
     powermeasure(μ, dims)
 end
 
-Base.:^(μ::AbstractMeasure, dims::Tuple) = powermeasure(μ, one_to.(dims))
+Base.:^(μ::AbstractMeasure, dims::Tuple) = powermeasure(μ, maybestatic_oneto.(dims))
 Base.:^(μ::AbstractMeasure, n) = powermeasure(μ, (n,))
 
 # Base.show(io::IO, d::PowerMeasure) = print(io, d.parent, " ^ ", size(d.xs))
@@ -117,7 +118,10 @@ for func in [:logdensityof, :logdensity_def]
         end
     end
 
-    @eval @inline function $func(d::PowerMeasure{<:Any,Tuple{<:StaticOneToLike}}, x)
+    @eval @inline function $func(
+        d::PowerMeasure{<:Any,Tuple{<:StaticOneToLike{N}}},
+        x,
+    ) where {N}
         parent = d.parent
         sum(1:N) do j
             @inbounds $func(parent, x[j])
@@ -171,8 +175,7 @@ logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0)
 
 # To avoid ambiguities
 function logdensity_def(
-    ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0},N}}},
-    x,
+    ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0},N}}}, ::Any,
 ) where {P<:PrimitiveMeasure,N}
     static(0.0)
 end
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index 4b957651..e7244fac 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -13,7 +13,7 @@ function transport_def(ν::StdMeasure, μ::PowerMeasure{<:StdMeasure}, x)
 end
 
 function transport_def(ν::PowerMeasure{<:StdMeasure}, μ::StdMeasure, x)
-    return fill_with(transport_def(ν.parent, μ, only(x)), map(length, ν.axes))
+    return maybestatic_fill(transport_def(ν.parent, μ, only(x)), map(length, ν.axes))
 end
 
 function transport_def(
diff --git a/src/static.jl b/src/static.jl
deleted file mode 100644
index 12db9585..00000000
--- a/src/static.jl
+++ /dev/null
@@ -1,365 +0,0 @@
-# A lots of this is about bridging Static and StaticArrays, both have their
-# own SUnitRange and SOneTo. Also provides tools to control static vs dynamic
-# array, size and axes handling.
-
-"""
-    MeasureBase.StaticUnitRange
-
-The MeasureBase default type for static unit ranges.
-"""
-const StaticUnitRange = @static if isdefined(StaticArrays, :SUnitRange)
-    # Unclear if StaticArrays.SUnitRange is part of StaticArrays stable API.
-    # Some packages use it, but let's be careful in case it disappears.
-    StaticArrays.SUnitRange
-else
-    Static.SUnitRange
-end
-
-"""
-    MeasureBase.StaticOneTo
-
-The MeasureBase default type for static one-based unit ranges.
-"""
-const StaticOneTo{T} = StaticArrays.SOneTo{T}
-
-"""
-    MeasureBase.IntegerLike
-
-Equivalent to `Union{Integer,Static.StaticInteger}`.
-"""
-const IntegerLike = Union{Integer,Static.StaticInteger}
-
-"""
-    MeasureBase.SizeLike
-
-Something that can represent the size of a collection.
-"""
-const SizeLike = Union{Tuple{},Tuple{Vararg{IntegerLike}},StaticArrays.Size}
-
-"""
-    MeasureBase.StaticSizeLike
-
-Something that can represent the size of a statically sized collection.
-"""
-const StaticSizeLike = Union{Tuple{Vararg{StaticInteger}},StaticArrays.Size}
-
-"""
-    MeasureBase.AxesLike
-
-Something that can represent axes of a collection.
-"""
-const AxesLike = Union{Tuple{},Tuple{Vararg{AbstractVector{<:IntegerLike}}}}
-
-"""
-    MeasureBase.StaticAxesLike
-
-Something that can represent axes of a statically sized collection.
-"""
-@static if isdefined(StaticArrays, :SUnitRange)
-    const StaticAxesLike = Union{
-        Tuple{Vararg{Union{StaticArrays.SOneTo,StaticArrays.SUnitRange,Static.SUnitRange}}},
-    }
-else
-    const StaticAxesLike =
-        Union{Tuple{Vararg{Union{StaticArrays.SOneTo,Static.SUnitRange}}}}
-end
-
-"""
-    const OneToLike
-
-Alias for unit ranges that start at one.
-"""
-const OneToLike = Union{Base.OneTo,StaticArrays.SOneTo,Static.SOneTo}
-
-"""
-    const StaticOneToLike{N}
-
-A static unit range from one to N.
-"""
-const StaticOneToLike{N} = Union{StaticArrays.SOneTo{N},Static.SOneTo{N}}
-
-"""
-    const StaticUnitRangeLike
-
-A static unit range.
-"""
-@static if isdefined(StaticArrays, :SUnitRange)
-    const StaticUnitRangeLike =
-        Union{StaticArrays.SOneTo,StaticArrays.SUnitRange,Static.SUnitRange}
-else
-    const StaticUnitRangeLike = Union{StaticArrays.SOneTo,Static.SUnitRange}
-end
-
-"""
-    MeasureBase.one_to(n::IntegerLike)
-
-Creates a range from one to n.
-
-Returns an instance of `Base.OneTo` or `Static.SOneTo`, depending
-on the type of `n`.
-"""
-@inline one_to(n::Integer) = Base.OneTo(n)
-@inline one_to(::Static.StaticInteger{N}) where {N} = Static.SOneTo{N}()
-
-"""
-    MeasureBase.asnonstatic(x)
-
-Return a non-static equivalent of `x`.
-
-Defaults to `Static.dynamic(x)`.
-"""
-@inline asnonstatic(x::Number) = dynamic(x)
-@inline asnonstatic(::Tuple{}) = ()
-@static if isdefined(StaticArrays, :SUnitRange)
-    @inline asnonstatic(r::StaticArrays.SUnitRange) = r[begin]:r[end]
-end
-@inline asnonstatic(r::AbstractUnitRange) = asnonstatic(r[begin]):asnonstatic(r[end])
-@inline asnonstatic(r::Base.OneTo) = Base.OneTo(asnonstatic(r.stop))
-@inline asnonstatic(::StaticOneToLike{N}) where {N} = Base.OneTo(N)
-@inline asnonstatic(x::SizeLike) = map(asnonstatic, x)
-@inline asnonstatic(::StaticArrays.Size{TPL}) where {TPL} = TPL
-@inline asnonstatic(x::AxesLike) = map(asnonstatic, x)
-
-"""
-    MeasureBase.fill_with(x, sz::NTuple{N,<:IntegerLike}) where N
-
-Creates an array of size `sz` filled with `x`.
-
-The result will typically be either a `FillArrays.Fill` or a static array,
-"""
-function fill_with end
-
-@inline fill_with(x::T, n::IntegerLike) where {T} = fill_with(x, (n,))
-
-@inline fill_with(x::T, ::Tuple{}) where {T} = FillArrays.Fill(x)
-
-@inline fill_with(x, sz::SizeLike) = fill_with(x, size2axes(sz))
-
-@inline function fill_with(x::T, sz::StaticSizeLike) where {T}
-    fill(x, staticarray_type(T, canonical_size(sz)))
-end
-
-@inline function fill_with(x, axs::AxesLike)
-    dyn_axs = map(asnonstatic, axs)
-    FillArrays.Fill(x, dyn_axs)
-end
-
-# While `FillArrays.Fill` (mostly?) works with axes that are static unit
-# ranges, some operations that automatic differentiation requires do fail
-# on such instances of `Fill` (e.g. `reshape` from dynamic to static size).
-# So need to build a filled static array:
-@inline function fill_with(x::T, axs::Tuple{Vararg{StaticOneToLike}}) where {T}
-    sz = axes2size(axs)
-    fill(x, staticarray_type(T, sz))
-end
-
-"""
-    MeasureBase.staticarray_type(T, sz::StaticArrays.Size)
-
-Returns the type of a static array with element type `T` and size `sz`.
-"""
-function staticarray_type end
-
-@inline @generated function staticarray_type(
-    ::Type{T},
-    ::StaticArrays.Size{sz},
-) where {T,sz}
-    N = length(sz)
-    len = prod(sz)
-    :(SArray{Tuple{$sz...},T,$N,$len})
-end
-
-"""
-    MeasureBase.maybestatic_reshape(A, sz)
-
-Reshapes array `A` to sizes `sz`.
-
-If `A` is a static array and `sz` is static, the result is a static array.
-"""
-function maybestatic_reshape end
-
-maybestatic_reshape(A, sz) = reshape(A, canonical_size(sz))
-function maybestatic_reshape(A, sz::StaticSizeLike)
-    StaticArrays.SArray(reshape(A, canonical_size(sz)))
-end
-function maybestatic_reshape(A::StaticArray, sz::Tuple{Vararg{StaticInteger}})
-    staticarray_type(eltype(A), canonical_size(sz))(Tuple(A))
-end
-
-"""
-    MeasureBase.maybestatic_length(x)
-
-Returns the length of `x` as a dynamic or static integer.
-"""
-@inline maybestatic_length(::Number) = static(1)
-@inline maybestatic_length(::Tuple{}) = static(0)
-@inline maybestatic_length(::Tuple{Vararg{Any,N}}) where {N} = static(N)
-@inline maybestatic_length(nt::NamedTuple) = maybestatic_length(values(nt))
-@inline maybestatic_length(A::AbstractArray) = size2length(maybestatic_size(A))
-@static if isdefined(StaticArrays, :SUnitRange)
-    @inline maybestatic_length(r::StaticArrays.SUnitRange) =
-        maybestatic_last(r) - maybestatic_first(r) + static(1)
-end
-@inline maybestatic_length(r::AbstractUnitRange) =
-    maybestatic_last(r) - maybestatic_first(r) + static(1)
-@inline maybestatic_length(r::Base.OneTo) = length(r)
-@inline maybestatic_length(::StaticArrays.SOneTo{N}) where {N} = static(N)
-@inline maybestatic_length(::Static.SOneTo{N}) where {N} = static(N)
-
-"""
-
-    MeasureBase.maybestatic_size(x)
-
-Returns the size of `x` as a tuple of dynamic or static integers.
-"""
-@inline maybestatic_size(::Number) = ()
-@inline maybestatic_size(::Tuple{}) =
-    throw(ArgumentError("Cannot determine (maybe-static) size of empty tuple"))
-@inline maybestatic_size(::Tuple{Vararg{Any,N}}) where {N} = StaticArrays.Size{(N,)}()
-@inline maybestatic_size(nt::NamedTuple) = maybestatic_size(values(nt))
-@inline maybestatic_size(A::AbstractArray) = axes2size(maybestatic_axes(A))
-@inline maybestatic_size(A::StaticArray) = StaticArrays.Size(A)
-
-"""
-    MeasureBase.maybestatic_axes(x)::Tuple{Vararg{IntegerLike}}
-
-Returns the size of `x` as a tuple of dynamic or static integers.
-"""
-@inline maybestatic_axes(::Number) = ()
-
-@inline maybestatic_axes(::Tuple{}) = (StaticOneTo(0),)
-@inline maybestatic_axes(::Tuple{Vararg{Any,N}}) where {N} = (StaticOneTo(N),)
-@inline maybestatic_axes(nt::NamedTuple) = maybestatic_axes(values(nt))
-@inline maybestatic_axes(::StaticOneToLike{N}) where {N} = (StaticOneTo(N),)
-@static if isdefined(StaticArrays, :SUnitRange)
-    @inline maybestatic_axes(r::StaticArrays.SUnitRange) = axes(r)
-end
-@inline maybestatic_axes(r::Static.OptionallyStaticUnitRange) = canonical_axes(axes(r))
-@inline maybestatic_axes(r::AbstractUnitRange) = axes(r)
-@inline maybestatic_axes(A::AbstractArray) = axes(A)
-@inline maybestatic_axes(A::StaticArray) = axes(A)
-
-"""
-    MeasureBase.axes2size(x::Tuple)
-    MeasureBase.axes2size(x::StaticArrays.Size)
-
-Get a length from a size (tuple).
-"""
-@inline axes2size(::Tuple{}) = ()
-@inline axes2size(axs::Tuple) = canonical_size(map(maybestatic_length, axs))
-
-"""map(maybestatic_length, axs)
-    MeasureBase.size2axes(sz::Tuple)
-    MeasureBase.size2axes(sz::StaticArrays.Size)
-
-Get one-based indexing axes from a size.
-"""
-@inline size2axes(::Tuple{}) = ()
-@inline size2axes(sz::Tuple) = canonical_axes(map(one_to, sz))
-@inline size2axes(::StaticArrays.Size{TPL}) where {TPL} = map(StaticOneTo, TPL)
-
-"""
-    MeasureBase.size2length(sz::Tuple)
-    MeasureBase.size2length(sz::StaticArrays.Size)
-
-Get a length from a size (tuple).
-"""
-@inline size2length(::Tuple{}) = static(1)
-@inline size2length(sz::Tuple) = prod(sz)
-@inline size2length(::StaticArrays.Size{TPL}) where {TPL} = static(prod(TPL))
-
-"""
-    MeasureBase.asaxes(axs::AxesLike)
-    MeasureBase.asaxes(sz::SizeLike)
-    MeasureBase.asaxes(len::IntegerLike)
-
-Converts axes or a size or a length of a collection to axes.
-
-One-based indexing will be used if the indexing offset can't be inferred from
-the given dimensions.
-"""
-@inline asaxes(::Tuple{}) = ()
-@inline asaxes(axs::AxesLike) = axs
-@inline asaxes(sz::SizeLike) = size2axes(sz)
-@inline asaxes(len::IntegerLike) = size2axes((len,))
-
-"""
-    MeasureBase.maybestatic_eachindex(x)
-
-Returns the the index range of `x` as a dynamic or static integer range
-"""
-maybestatic_eachindex(::Tuple{}) = StaticOneTo(0)
-maybestatic_eachindex(::Tuple{Vararg{Any,N}}) where {N} = StaticOneTo(N)
-maybestatic_eachindex(nt::NamedTuple) = maybestatic_eachindex(values(nt))
-maybestatic_eachindex(x::AbstractArray) = canonical_indices(eachindex(x))
-
-"""
-    MeasureBase.maybestatic_first(A)
-
-Returns the first element of `A` as a dynamic or static value.
-"""
-maybestatic_first(tpl::Tuple) = tpl[begin]
-maybestatic_first(nt::NamedTuple) = nt[begin]
-maybestatic_first(A::AbstractArray) = A[begin]
-maybestatic_first(::StaticArrays.Size{tpl}) where {tpl} = static(tpl[begin])
-maybestatic_first(::StaticArrays.SOneTo{N}) where {N} = static(1)
-@static if isdefined(StaticArrays, :SUnitRange)
-    maybestatic_first(::StaticArrays.SUnitRange{B,L}) where {B,L} = static(B)
-end
-function maybestatic_first(
-    ::Static.OptionallyStaticUnitRange{<:Static.StaticInteger{from},<:Static.StaticInteger},
-) where {from}
-    static(from)
-end
-
-"""
-    MeasureBase.maybestatic_last(A)
-
-Returns the last element of `A` as a dynamic or static value.
-"""
-maybestatic_last(tpl::Tuple) = tpl[end]
-maybestatic_last(nt::NamedTuple) = nt[end]
-maybestatic_last(A::AbstractArray) = A[end]
-maybestatic_last(::StaticArrays.Size{tpl}) where {tpl} = static(tpl[end])
-maybestatic_last(::StaticArrays.SOneTo{N}) where {N} = static(N)
-@static if isdefined(StaticArrays, :SUnitRange)
-    maybestatic_last(::StaticArrays.SUnitRange{B,L}) where {B,L} = static(B + L - 1)
-end
-function maybestatic_last(
-    ::Static.OptionallyStaticUnitRange{<:Any,<:Static.StaticInteger{until}},
-) where {until}
-    static(until)
-end
-
-"""
-    MeasureBase.canonical_indices(idxs::AbstractVector{<:IntegerLike})
-
-Return the canonical representation of a collection axis indices.
-"""
-@inline canonical_indices(idxs::AbstractVector{<:IntegerLike}) = idxs
-@inline canonical_indices(idxs::AbstractArray{<:CartesianIndex}) = idxs
-@inline canonical_indices(
-    ::Static.OptionallyStaticUnitRange{<:StaticInteger{1},<:StaticInteger{N}},
-) where {N} = StaticArrays.SOneTo{N}()
-@inline canonical_indices(
-    ::Static.OptionallyStaticUnitRange{<:StaticInteger{A},<:StaticInteger{B}},
-) where {A,B} = StaticUnitRange(A, B)
-@inline canonical_indices(
-    r::Static.OptionallyStaticUnitRange{<:StaticInteger{1},<:Integer},
-) = Base.OneTo(last(r))
-
-"""
-    MeasureBase.canonical_size(sz::SizeLike)
-
-Return the canonical representation of a collection size.
-"""
-@inline canonical_size(sz::SizeLike) = sz
-@inline canonical_size(sz::Tuple{Vararg{Static.StaticInteger}}) =
-    StaticArrays.Size{map(dynamic, sz)}()
-
-"""
-    MeasureBase.canonical_axes(sz::SizeLike)
-
-Return the canonical representation collection axes.
-"""
-@inline canonical_axes(axs::AxesLike) = map(canonical_indices, axs)
diff --git a/test/static.jl b/test/static.jl
deleted file mode 100644
index 83092ec2..00000000
--- a/test/static.jl
+++ /dev/null
@@ -1,325 +0,0 @@
-using Test
-
-import MeasureBase
-using MeasureBase:
-    StaticUnitRange,
-    StaticOneTo,
-    IntegerLike,
-    SizeLike,
-    StaticSizeLike,
-    AxesLike,
-    StaticAxesLike,
-    OneToLike,
-    StaticOneToLike,
-    StaticUnitRangeLike,
-    one_to,
-    asnonstatic,
-    fill_with,
-    staticarray_type,
-    maybestatic_reshape,
-    maybestatic_length,
-    maybestatic_size,
-    maybestatic_axes,
-    axes2size,
-    size2axes,
-    size2length,
-    asaxes,
-    maybestatic_eachindex,
-    maybestatic_first,
-    maybestatic_last,
-    canonical_indices,
-    canonical_size,
-    canonical_axes
-
-import Static
-using Static: static
-import StaticArrays
-import FillArrays
-
-@testset "static" begin
-    v = 4.2
-    T = typeof(v)
-
-    tpl = (7, 42, 5)
-    nt = (a = 7, b = 42, c = 5)
-
-    i = 7
-    si = static(7)
-
-    @test i isa IntegerLike
-    @test si isa IntegerLike
-
-    sz = (2, 4, 3)
-    sasz = StaticArrays.Size(2, 4, 3)
-    sisz = (static(2), static(4), static(3))
-
-    len = prod(sz)
-    slen = static(len)
-
-    @test sz isa SizeLike
-    @test sasz isa SizeLike
-    @test sisz isa SizeLike
-
-    @test !(sz isa StaticSizeLike)
-    @test sasz isa StaticSizeLike
-    @test sisz isa StaticSizeLike
-
-    axs = (Base.OneTo(2), 2:5, Base.OneTo(3))
-    axs1 = (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3))
-    saaxs = (StaticOneTo(2), StaticUnitRange(2, 5), StaticOneTo(3))
-    saaxs1 = (StaticOneTo(2), StaticOneTo(4), StaticOneTo(3))
-    siaxs = (Static.SOneTo(2), static(2):static(5), static(1):static(3))
-    siaxs1 = (Static.SOneTo(2), static(1):static(4), static(1):static(3))
-
-    @test axs isa AxesLike
-    @test axs1 isa AxesLike
-    @test saaxs isa AxesLike
-    @test saaxs1 isa AxesLike
-    @test siaxs isa AxesLike
-    @test siaxs1 isa AxesLike
-
-    @test !(axs isa StaticAxesLike)
-    @test saaxs isa StaticAxesLike
-    @test saaxs1 isa StaticAxesLike
-    @test siaxs isa StaticAxesLike
-    @test siaxs1 isa StaticAxesLike
-
-    @test axs[1] isa OneToLike
-    @test !(axs[2] isa OneToLike)
-    @test axs[3] isa OneToLike
-
-    @test saaxs[1] isa OneToLike
-    @test !(saaxs[2] isa OneToLike)
-    @test saaxs1[2] isa OneToLike
-    @test saaxs[3] isa OneToLike
-
-    @test siaxs[1] isa OneToLike
-    @test !(siaxs[2] isa OneToLike)
-    @test siaxs1[2] isa OneToLike
-    @test siaxs[3] isa OneToLike
-
-    @test !(axs[1] isa StaticOneToLike)
-    @test !(axs[2] isa StaticOneToLike)
-    @test !(axs[3] isa StaticOneToLike)
-
-    @test saaxs[1] isa StaticOneToLike
-    @test !(saaxs[2] isa StaticOneToLike)
-    @test saaxs1[2] isa StaticOneToLike
-    @test saaxs[3] isa StaticOneToLike
-
-    @test siaxs[1] isa StaticOneToLike
-    @test !(siaxs[2] isa StaticOneToLike)
-    @test siaxs1[2] isa StaticOneToLike
-    @test siaxs[3] isa StaticOneToLike
-
-    @test !(axs[1] isa StaticUnitRangeLike)
-    @test !(axs[2] isa StaticUnitRangeLike)
-    @test !(axs[3] isa StaticUnitRangeLike)
-
-    @test saaxs[1] isa StaticUnitRangeLike
-    @test saaxs[2] isa StaticUnitRangeLike
-    @test saaxs1[2] isa StaticUnitRangeLike
-    @test saaxs[3] isa StaticUnitRangeLike
-
-    @test siaxs[1] isa StaticUnitRangeLike
-    @test siaxs[2] isa StaticUnitRangeLike
-    @test siaxs1[2] isa StaticUnitRangeLike
-    @test siaxs[3] isa StaticUnitRangeLike
-
-    @test @inferred(one_to(i)) == Base.OneTo(i)
-    @test @inferred(one_to(si)) == StaticOneTo(i)
-
-    @test @inferred(asnonstatic(i)) === i
-    @test @inferred(asnonstatic(si)) === i
-    @test @inferred(asnonstatic(sz)) === sz
-    @test @inferred(asnonstatic(sasz)) === sz
-    @test @inferred(asnonstatic(sisz)) === sz
-    @test @inferred(asnonstatic(axs)) === axs
-    @test @inferred(asnonstatic(saaxs)) === axs
-    @test @inferred(asnonstatic(saaxs1)) === (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3))
-    @test @inferred(asnonstatic(siaxs)) === axs
-    @test @inferred(asnonstatic(siaxs1)) === (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3))
-
-    @test @inferred(fill_with(v, i)) === FillArrays.Fill(v, i)
-    @test @inferred(fill_with(v, si)) === StaticArrays.SVector(fill(v, i)...)
-    @test @inferred(fill_with(v, ())) === FillArrays.Fill(v)
-
-    @test @inferred(fill_with(v, sz)) === FillArrays.Fill(v, sz)
-    @test @inferred(fill_with(v, sasz)) === StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz))
-    @test @inferred(fill_with(v, sisz)) === StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz))
-
-    @test @inferred(fill_with(v, axs)) === FillArrays.Fill(v, axs)
-    @test @inferred(fill_with(v, saaxs)) === FillArrays.Fill(v, axs)
-    @test @inferred(fill_with(v, saaxs1)) ===
-          StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz))
-    @test @inferred(fill_with(v, siaxs)) === FillArrays.Fill(v, axs)
-    @test @inferred(fill_with(v, siaxs1)) ===
-          StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz))
-
-    @test @inferred(staticarray_type(T, sasz)) <: StaticArrays.SArray{Tuple{2,4,3},T}
-
-    A = rand(T, len)
-    FA = FillArrays.Fill(v, len)
-    SA = StaticArrays.SVector(A...)
-
-    # Array with CartesianIndices
-    ciA = view(rand(5, 6, 6), 3:4, 2:5, 3:5)
-    ciidxs = eachindex(ciA)
-
-    rshpA = reshape(A, sz)
-    rshpFA = FillArrays.Fill(v, sz)
-    rshpSA = StaticArrays.SArray{Tuple{sz...},T}(A)
-
-    @test @inferred(maybestatic_reshape(A, sz)) == rshpA
-    @test typeof(maybestatic_reshape(A, sz)) == typeof(rshpA)
-    @test @inferred(maybestatic_reshape(A, sasz)) == rshpA
-    @test maybestatic_reshape(A, sasz) isa StaticArrays.SArray
-    @test @inferred(maybestatic_reshape(A, sisz)) == rshpA
-    @test maybestatic_reshape(A, sisz) isa StaticArrays.SArray
-
-    @test @inferred(maybestatic_reshape(FA, sz)) == rshpFA
-    @test typeof(maybestatic_reshape(FA, sz)) == typeof(rshpFA)
-    @test @inferred(maybestatic_reshape(FA, sasz)) == rshpFA
-    @test maybestatic_reshape(FA, sasz) isa StaticArrays.SArray
-    @test @inferred(maybestatic_reshape(FA, sisz)) == rshpFA
-    @test maybestatic_reshape(FA, sisz) isa StaticArrays.SArray
-
-    @test @inferred(maybestatic_reshape(SA, sz)) == rshpA
-    @test maybestatic_reshape(SA, sz) isa Base.ReshapedArray{T,3,<:StaticArrays.SVector}
-    @test @inferred(maybestatic_reshape(SA, sasz)) === rshpSA
-    @test @inferred(maybestatic_reshape(SA, sisz)) === rshpSA
-
-    @test @inferred(maybestatic_length(5)) === static(1)
-    @test @inferred(maybestatic_length(())) === static(0)
-    @test @inferred(maybestatic_length((sz))) === static(3)
-    @test @inferred(maybestatic_length((a = 2, b = 4, c = 3))) === static(3)
-    @test @inferred(maybestatic_length(Base.OneTo(4))) === 4
-    @test @inferred(maybestatic_length(StaticArrays.SOneTo(4))) === static(4)
-    @test @inferred(maybestatic_length(Static.SOneTo(4))) === static(4)
-    @test @inferred(maybestatic_length(static(2):static(5))) === static(4)
-    @test @inferred(maybestatic_length(rshpA)) === length(rshpA)
-    @test @inferred(maybestatic_length(rshpFA)) === length(rshpA)
-    @test @inferred(maybestatic_length(rshpSA)) === static(length(rshpA))
-
-    @test @inferred(maybestatic_size(5)) === ()
-    @test_throws ArgumentError maybestatic_size(())
-    @test @inferred(maybestatic_size((sz))) === StaticArrays.Size(3)
-    @test @inferred(maybestatic_size((a = 2, b = 4, c = 3))) === StaticArrays.Size(3)
-    @test @inferred(maybestatic_size(Base.OneTo(4))) === (4,)
-    @test @inferred(maybestatic_size(StaticArrays.SOneTo(4))) === StaticArrays.Size(4)
-    @test @inferred(maybestatic_size(StaticUnitRange(2, 5))) === StaticArrays.Size(4)
-    @test @inferred(maybestatic_size(Static.SOneTo(4))) === StaticArrays.Size(4)
-    @test @inferred(maybestatic_size(static(2):static(5))) === StaticArrays.Size(4)
-    @test @inferred(maybestatic_size(rshpA)) === size(rshpA)
-    @test @inferred(maybestatic_size(rshpFA)) === size(rshpA)
-    @test @inferred(maybestatic_size(rshpSA)) === StaticArrays.Size(size(rshpA)...)
-
-    @test @inferred(maybestatic_axes(5)) === ()
-    @test @inferred(maybestatic_axes(())) === (StaticOneTo(0),)
-    @test @inferred(maybestatic_axes((sz))) === (StaticOneTo(3),)
-    @test @inferred(maybestatic_axes((a = 2, b = 4, c = 3))) === (StaticOneTo(3),)
-    @test @inferred(maybestatic_axes(Base.OneTo(4))) === (Base.OneTo(4),)
-    @test @inferred(maybestatic_axes(StaticArrays.SOneTo(4))) === (StaticOneTo(4),)
-    @test @inferred(maybestatic_axes(Static.SOneTo(4))) === (StaticOneTo(4),)
-    @test @inferred(maybestatic_axes(static(2):static(5))) === (StaticOneTo(4),)
-    @test @inferred(maybestatic_axes(rshpA)) === axes(rshpA)
-    @test @inferred(maybestatic_axes(rshpFA)) === axes(rshpA)
-    @test @inferred(maybestatic_axes(rshpSA)) === saaxs1
-
-    @test @inferred(axes2size(())) === ()
-    @test @inferred(axes2size(axs)) === sz
-    @test @inferred(axes2size(saaxs)) === sasz
-    @test @inferred(axes2size(saaxs1)) === sasz
-    @test @inferred(axes2size(siaxs)) === sasz
-    @test @inferred(axes2size(siaxs1)) === sasz
-
-    @test @inferred(size2axes(())) === ()
-    @test @inferred(size2axes(sz)) === axs1
-    @test @inferred(size2axes(sasz)) === saaxs1
-    @test @inferred(size2axes(sisz)) === saaxs1
-
-    @test @inferred(size2length(())) === static(1)
-    @test @inferred(size2length(sz)) === len
-    @test @inferred(size2length(sasz)) === slen
-    @test @inferred(size2length(sisz)) === slen
-
-    @test @inferred(asaxes(())) === ()
-    @test @inferred(asaxes(len)) === (Base.OneTo(len),)
-    @test @inferred(asaxes(slen)) === (StaticOneTo(len),)
-    @test @inferred(asaxes(sz)) === axs1
-    @test @inferred(asaxes(sasz)) === saaxs1
-    @test @inferred(asaxes(sisz)) === saaxs1
-    @test @inferred(asaxes(axs)) === axs
-    @test @inferred(asaxes(axs1)) === axs1
-    @test @inferred(asaxes(saaxs)) === saaxs
-    @test @inferred(asaxes(saaxs1)) === saaxs1
-    @test @inferred(asaxes(siaxs)) === siaxs
-    @test @inferred(asaxes(siaxs1)) === siaxs1
-
-    @test @inferred(maybestatic_eachindex(())) === StaticOneTo(0)
-    @test @inferred(maybestatic_eachindex(tpl)) === StaticOneTo(3)
-    @test @inferred(maybestatic_eachindex(nt)) === StaticOneTo(3)
-    @test @inferred(maybestatic_eachindex(axs[1])) === Base.OneTo(length(axs[1]))
-    @test @inferred(maybestatic_eachindex(axs[2])) === Base.OneTo(length(axs[2]))
-    @test @inferred(maybestatic_eachindex(saaxs[1])) === StaticOneTo(length(axs[1]))
-    @test @inferred(maybestatic_eachindex(saaxs[2])) === StaticOneTo(length(axs[2]))
-    @test @inferred(maybestatic_eachindex(siaxs[1])) === StaticOneTo(length(axs[1]))
-    @test @inferred(maybestatic_eachindex(siaxs[2])) === StaticOneTo(length(axs[2]))
-    @test @inferred(maybestatic_eachindex(A)) === Base.OneTo(24)
-    @test @inferred(maybestatic_eachindex(ciA)) === eachindex(ciA)
-    @test @inferred(maybestatic_eachindex(FA)) === Base.OneTo(24)
-    @test @inferred(maybestatic_eachindex(SA)) === StaticOneTo(24)
-
-    @test_throws BoundsError maybestatic_first(())
-    @test @inferred(maybestatic_first(tpl)) === first(tpl)
-    @test @inferred(maybestatic_first(nt)) === first(nt)
-    @test @inferred(maybestatic_first(sz)) === first(sz)
-    @test @inferred(maybestatic_first(sasz)) === static(first(sz))
-    @test @inferred(maybestatic_first(sisz)) === static(first(sz))
-    @test @inferred(maybestatic_first(axs[1])) === first(axs[1])
-    @test @inferred(maybestatic_first(axs[2])) === first(axs[2])
-    @test @inferred(maybestatic_first(saaxs[1])) === static(first(axs[1]))
-    @test @inferred(maybestatic_first(saaxs[2])) === static(first(axs[2]))
-    @test @inferred(maybestatic_first(siaxs[1])) === static(first(axs[1]))
-    @test @inferred(maybestatic_first(siaxs[2])) === static(first(axs[2]))
-    @test @inferred(maybestatic_first(A)) === first(A)
-    @test @inferred(maybestatic_first(ciA)) === first(ciA)
-    @test @inferred(maybestatic_first(FA)) === first(FA)
-    @test @inferred(maybestatic_first(SA)) === first(SA)
-
-    @test_throws BoundsError maybestatic_last(())
-    @test @inferred(maybestatic_last(tpl)) === last(tpl)
-    @test @inferred(maybestatic_last(nt)) === last(nt)
-    @test @inferred(maybestatic_last(sz)) === last(sz)
-    @test @inferred(maybestatic_last(sasz)) === static(last(sz))
-    @test @inferred(maybestatic_last(sisz)) === static(last(sz))
-    @test @inferred(maybestatic_last(axs[1])) === last(axs[1])
-    @test @inferred(maybestatic_last(axs[2])) === last(axs[2])
-    @test @inferred(maybestatic_last(saaxs[1])) === static(last(axs[1]))
-    @test @inferred(maybestatic_last(saaxs[2])) === static(last(axs[2]))
-    @test @inferred(maybestatic_last(siaxs[1])) === static(last(axs[1]))
-    @test @inferred(maybestatic_last(siaxs[2])) === static(last(axs[2]))
-    @test @inferred(maybestatic_last(A)) === last(A)
-    @test @inferred(maybestatic_last(ciA)) === last(ciA)
-    @test @inferred(maybestatic_last(FA)) === last(FA)
-    @test @inferred(maybestatic_last(SA)) === last(SA)
-
-    @test @inferred(canonical_indices(axs[1])) === axs[1]
-    @test @inferred(canonical_indices(axs[2])) === axs[2]
-    @test @inferred(canonical_indices(saaxs[1])) === saaxs[1]
-    @test @inferred(canonical_indices(saaxs[2])) === saaxs[2]
-    @test @inferred(canonical_indices(siaxs[1])) === saaxs[1]
-    @test @inferred(canonical_indices(siaxs[2])) === saaxs[2]
-    @test @inferred(canonical_indices(ciidxs)) === ciidxs
-
-    @test @inferred(canonical_size(sz)) === sz
-    @test @inferred(canonical_size(sasz)) === sasz
-    @test @inferred(canonical_size(sisz)) === sasz
-
-    @test @inferred(canonical_axes(axs)) === axs
-    @test @inferred(canonical_axes(axs1)) === axs1
-    @test @inferred(canonical_axes(saaxs)) === saaxs
-    @test @inferred(canonical_axes(saaxs1)) === saaxs1
-    @test @inferred(canonical_axes(siaxs)) === saaxs
-    @test @inferred(canonical_axes(siaxs1)) === saaxs1
-end

From c436e4b4dd599b068f652f4f7fdb7d496b7a17df Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:10 +0200
Subject: [PATCH 016/122] Remove ZeroSet and CodimOne

Currently unused and undocumented, can add it back later when needed.
---
 src/domains.jl | 60 --------------------------------------------------
 1 file changed, 60 deletions(-)

diff --git a/src/domains.jl b/src/domains.jl
index e03f753c..c9912420 100644
--- a/src/domains.jl
+++ b/src/domains.jl
@@ -69,77 +69,17 @@ function Base.getindex(::typeof(ℤ), r::AbstractUnitRange)
     BoundedInts(extrema(r)...)
 end
 
-###########################################################
-# ZeroSet
 
-export ZeroSet
-
-struct ZeroSet{F,G} <: AbstractDomain
-    f::F
-    ∇f::G
-end
-
-# Based on some quick tests, but may need some adjustment
-Base.in(x::AbstractArray{T}, z::ZeroSet) where {T} = abs(z.f(x)) < ldexp(eps(float(T)), 6)
-
-###########################################################
-# CodimOne
-
-export CodimOne
-
-abstract type CodimOne <: AbstractDomain end
-
-function tangentat(
-    a::CodimOne,
-    b::CodimOne,
-    x::AbstractArray{T};
-    tol = ldexp(eps(float(T)), 6),
-) where {T}
-    # Sometimes you get lucky
-    a == b && return true
-
-    # Get the normal vectors
-    g1 = a.∇f(x)
-    g2 = b.∇f(x)
-
-    # See if one is a multiple of the other
-    one(T) - Statistics.corm(g1, zero(T), g2, zero(T)) < tol
-end
-
-function zeroset(::CodimOne)::ZeroSet end
-
-###########################################################
-# Simplex
-export Simplex
-
-struct Simplex <: CodimOne end
-
-function zeroset(::Simplex)
-    f(x::AbstractArray{T}) where {T} = sum(x) - one(T)
-    ∇f(x::AbstractArray{T}) where {T} = fill_with(one(T), size(x))
-    ZeroSet(f, ∇f)
-end
 
 function Base.in(x::AbstractArray{T}, ::Simplex) where {T}
     all(≥(zero(eltype(x))), x) || return false
     return x ∈ zeroset(Simplex())
 end
 
-projectto!(x, ::Simplex) = normalize!(x, 1)
 
-###########################################################
-# Sphere
 
 struct Sphere <: CodimOne end
 
-function zeroset(::Sphere)
-    f(x::AbstractArray{T}) where {T} = dot(x, x) - one(T)
-    ∇f(x::AbstractArray{T}) where {T} = x
-    ZeroSet(f, ∇f)
-end
-
 function Base.in(x::AbstractArray{T}, ::Sphere) where {T}
     return x ∈ zeroset(Sphere())
 end
-
-projectto!(x, ::Sphere) = normalize!(x, 2)

From 33b2407876436ae718107eedfde8ae49097ba175 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:10 +0200
Subject: [PATCH 017/122] Re-design of domains

---
 src/domains.jl             | 327 +++++++++++++++++++++++++++++++------
 src/primitives/counting.jl |   5 +
 src/primitives/lebesgue.jl |  18 +-
 3 files changed, 293 insertions(+), 57 deletions(-)

diff --git a/src/domains.jl b/src/domains.jl
index c9912420..c69cc29e 100644
--- a/src/domains.jl
+++ b/src/domains.jl
@@ -1,85 +1,316 @@
-abstract type AbstractDomain end
+"""
+    mdomain(m)::MeasureBase.SetLike
 
-abstract type RealDomain <: AbstractDomain end
+Return the domain, i.e. the measurable set, of the measure `m`.
 
-# TODO: Use IntervalSets
-struct RealNumbers <: RealDomain end
+The measure must allow for evaluating densities and the like over the whole
+domain, even if the support of the measure is only a subset of the domain.
 
-const ℝ = RealNumbers()
+May return [`MeasureBase.ImplicitDomain`](@ref) if the domain cannot be
+computed (efficiently).
+"""
+function mdomain end
+export mdomain
 
-Base.minimum(::RealNumbers) = static(-Inf)
-Base.maximum(::RealNumbers) = static(Inf)
+@inline mdomain(m) = ImplicitDomain(m)
 
-Base.in(x, ::RealNumbers) = isreal(x)
 
-Base.show(io::IO, ::typeof(ℝ)) = print(io, "ℝ")
+# Custom abstract set type. Design reserve to be able to switch to
+#`Base.AbstractSet` or another set type hierarchy in the future:
+"""
+    MeasureBase.ValueSet
 
-struct BoundedReals{L,U} <: RealDomain
-    lower::L
-    upper::U
+Abstract type for some measurable sets.
+
+Not every measurable set needs to be a a subtype of
+`MeasureBase.ValueSet`.
+
+See also [`MeasureBase.SetLike`](@ref).
+"""
+abstract type ValueSet end
+
+"""
+    const MeasureBase.SetLike = Union{MeaureBase.ValueSet, Base.AbstractSet, IntervalSets.Domain}
+
+Any kind of (measurable) set.
+
+There needs to be an implicit sigma-algebra for subtypes of
+`MeasureBase.SetLike` to make them useable for measures. This can't easily be
+imposed via type constraints, though, to is is by-contract.
+"""
+const SetLike = Union{MeasureBase.ValueSet,Base.AbstractSet,IntervalSets.Domain}
+
+"""
+    valdomain(x)::MeasureBase.SetLike
+
+Return the domain of a given value.
+
+May return [`MeasureBase.UnknownDomain`](@ref) if no domain type is available
+that can represents values like `x`.
+"""
+function valdomain end
+export valdomain
+
+@inline valdomain(x) = UnknownDomain(x)
+
+"""
+    MeasureBase.maybe_in(x, s)
+
+Test if `x` may be a member of `s`.
+
+Defaults to `in(x, s)`, but may be specialized for certain types of `s`,
+e.g. for `s::MeasureBase.ImplicitDomain`.
+"""
+function maybe_in end
+
+maybe_in(x, s) = in(x, s)
+
+"""
+    struct MeasureBase.ImplicitDomain{M} <: MeasureBase.ValueSet
+
+Represents the domain (i.e. the measurable set) of a measure `m::M`.
+
+Constructors:
+
+```
+MeasureBase.ImplicitDomain(m)
+```
+
+Fields:
+
+* `m::M`: The measure.
+
+For many pushforward measures and similar, the measureable space can not be
+computed efficiently or at all. In such cases, [`mdomain(m)`](@ref) should
+return `ImplicitDomain(m)`.
+
+Does not support `Base.in(x, s::MeasureBase.ImplicitDomain)`, and
+`MeasureBase.maybe_in(x, s::MeasureBase.ImplicitDomain)` always return `true`
+(unless specialized for the measure type).
+"""
+struct ImplicitDomain{M} <: ValueSet
+    m::M
+end
+
+@inline Base.union(s::ImplicitDomain, others::ImplicitDomain...) =
+    ImplicitDomain(+(s.m, map(x -> x.m, others)...))
+
+function Base.in(@nospecialize(x), ::ImplicitDomain)
+    throw(ArgumentError("Cannot test if a value lies withing an implicit domain."))
+end
+
+maybe_in(@nospecialize(x), ::ImplicitDomain) = true
+
+function Base.isempty(::ImplicitDomain)
+    throw(ArgumentError("Can't test if an ImplicitDomain is empty"))
+end
+
+"""
+    struct MeasureBase.UnknownDomain{T} <: MeasureBase.ValueSet
+
+Represents the unknown domain of a value of type `T`.
+
+Constructors:
+
+```
+MeasureBase.UnknownDomain(x::T)
+```
+
+Does not support `Base.in(x, s::MeasureBase.UnknownDomain)`, and
+`MeasureBase.maybe_in(x, s::MeasureBase.UnknownDomain)` always return `true`
+(unless specialized for the measure type).
+
+`isempty` will always return false, `UnknownDomain` should only be created
+if a value of type `T` existed in the first place, which implies that the
+domain can not be empty.
+"""
+struct UnknownDomain{T} <: ValueSet end
+
+UnknownDomain(::T) where {T} = UnknownDomain{T}()
+
+Base.eltype(::UnknownDomain{T}) where {T} = T
+
+@inline Base.union(s::UnknownDomain, others::UnknownDomain...) =
+    UnknownDomain{promote_type(eltype(s), map(eltype, others)...)}()
+
+function Base.in(@nospecialize(x), ::UnknownDomain)
+    throw(ArgumentError("Cannot test if a value lies withing an unknown domain."))
 end
 
-Base.in(x, b::BoundedReals) = b.lower ≤ x ≤ b.upper
+maybe_in(@nospecialize(x), ::UnknownDomain) = true
+
+Base.isempty(::UnknownDomain) = false
+
+"""
+    RealInterval() isa MeasureBase.ValueSet
+
+The real numbers.
+"""
+struct RealValues <: ValueSet end
+
+@inline Base.in(x::Real, ::RealValues) = true
+@inline Base.in(x, ::RealValues) = isreal(x)
+
+@inline Base.isempty(::RealValues) = false
+
+@inline Base.union(s::RealValues, ::RealValues...) = s
+
+@inline Base.minimum(::RealValues) = static(-Inf)
+@inline Base.maximum(::RealValues) = static(Inf)
+
+"""
+    const MeasureBase.ℝ = RealValues()
+
+The set of all real numbers, see [`MeasureBase.RealValues`](@ref).
+"""
+const ℝ = RealValues()
+
+Base.show(io::IO, ::MIME"text/plain", ::RealValues) = print(io, "MeasureBase.ℝ")
+
+"""
+    MeasureBase.IntegerValues() isa MeasureBase.ValueSet
+"""
+struct IntegerValues <: ValueSet end
+
+@inline Base.in(x::Integer, ::IntegerValues) = true
+@inline Base.in(x, ::IntegerValues) = isinteger(x)
 
-export ℝ, ℝ₊, 𝕀, ℤ
+@inline Base.isempty(::IntegerValues) = false
 
-const ℝ₊ = BoundedReals(static(0.0), static(Inf))
-const 𝕀 = BoundedReals(static(0.0), static(1.0))
+@inline Base.union(s::IntegerValues, ::IntegerValues...) = s
 
-Base.minimum(b::BoundedReals) = b.lower
-Base.maximum(b::BoundedReals) = b.upper
+# # This could get tricky with mixed-precision code. Probably needs some
+# # special AbstractInteger infinity type (but custom AbstractInteger types
+# # may cause a lot of method invalidations, which is why Static.StaticInteger
+# # is not an AbstractInteger).
+# @inline Base.minimum(::RealValues) = static(typemax(Int64))
+# @inline Base.maximum(::RealValues) = static(typemin(Int64))
 
-Base.show(io::IO, ::typeof(ℝ₊)) = print(io, "ℝ₊")
-Base.show(io::IO, ::typeof(𝕀)) = print(io, "𝕀")
+"""
+    const ℤ = IntegerValues()
 
-testvalue(::Type{T}, ::typeof(ℝ)) where {T} = zero(T)
-testvalue(::Type{T}, ::typeof(ℝ₊)) where {T} = one(T)
-testvalue(::Type{T}, ::typeof(𝕀)) where {T} = one(T) / 2
+The set of all integers, see [`MeasureBase.IntegerValues`](@ref).
+"""
+const ℤ = IntegerValues()
 
-abstract type IntegerDomain <: AbstractDomain end
+Base.show(io::IO, ::MIME"text/plain", ::IntegerValues) = print(io, "MeasureBase.ℤ")
 
-struct IntegerNumbers <: IntegerDomain end
+"""
+    struct MeasureBase.AbstractCartSetProd <: ValueSet
 
-Base.in(x, ::IntegerNumbers) = isinteger(x)
+Supertype for cartesian products of sets.
+"""
+abstract type AbstractCartSetProd <: ValueSet end
 
-const ℤ = IntegerNumbers()
+"""
+    struct CartesianProduct <: AbstractCartSetProd
 
-Base.show(io::IO, ::typeof(ℤ)) = print(io, "ℤ")
+A cartesian product over a collection of sets.
 
-Base.minimum(::IntegerNumbers) = static(-Inf)
-Base.maximum(::IntegerNumbers) = static(Inf)
-struct BoundedInts{L,U} <: IntegerDomain
-    lower::L
-    upper::U
+Constructor:
+
+```julia
+prodset = CartesianProduct(sets)
+```
+
+`sets` may be a `Tuple`, `NamedTuple` or `AbstractArray` of sets/domains.
+"""
+struct CartesianProduct{S<:Union{Tuple,NamedTuple,AbstractArray}} <: AbstractCartSetProd
+    _sets::S
+end
+
+componentsets(s::CartesianProduct) = s._sets
+
+setcartprod(sets::AbstractArray{<:SetLike}) = CartesianProduct(sets)
+setcartprod(sets::Tuple{Vararg{SetLike}}) = CartesianProduct(sets)
+setcartprod(sets::NamedTuple{names,<:Tuple{Vararg{SetLike}}}) = CartesianProduct(sets)
+
+@inline Base.in(x::Tuple{}, s::CartesianProduct{Tuple{}}) = true
+@inline Base.in(x::Tuple{Vararg{Any,N}}, s::CartesianProduct{<:Tuple{Vararg{Any,N}}}) where {N} =
+    prod(map(in, x, componentsets(s)))::Bool
+@inline Base.in(x::NamedTuple{names}, s::CartesianProduct{<:NamedTuple{names}}) where {names} =
+    prod(map(in, values(x), values(componentsets(s))))::Bool
+# ToDo: Allow this?
+# Base.in(x::AbstractVector, s::CartesianProduct{<:Tuple}) = all(in.(x,componentsets(s)))::Bool
+function Base.in(
+    x::AbstractArray{<:Any,N},
+    s::CartesianProduct{<:AbstractArray{<:Any,N}},
+) where {N}
+    sets = componentsets(s)
+    isempty(x) && isempty(sets) ? true : all(in.(x, sets))::Bool
+end
+
+@inline Base.isempty(s::CartesianProduct) = all(!isempty, componentsets(s))
+
+@inline function Base.union(
+    s::CartesianProduct{<:Tuple{Vararg{Any,N}}},
+    others::CartesianProduct{<:Tuple{Vararg{Any,N}}}...,
+) where {N}
+    CartesianProduct(map(union, componentsets(s), map(componentsets, others)...))
+end
+
+@inline function Base.union(
+    s::CartesianProduct{<:NamedTuple{names}},
+    others::CartesianProduct{<:NamedTuple{names}}...,
+) where {names}
+    CartesianProduct(map(union, componentsets(s), map(componentsets, others)...))
 end
 
-Base.in(x, b::BoundedInts) = x ∈ ℤ && b.lower ≤ x ≤ b.upper
+function Base.union(
+    s::CartesianProduct{<:AbstractArray{<:Any,N}},
+    others::CartesianProduct{<:AbstractArray{<:Any,N}}...,
+) where {N}
+    CartesianProduct(union.(componentsets(s), map(componentsets, others)...))
+end
 
-Base.minimum(b::BoundedInts) = b.lower
-Base.maximum(b::BoundedInts) = b.upper
+"""
+    struct CartesianPower <: AbstractCartSetProd
 
-function Base.show(io::IO, b::BoundedInts)
-    io = IOContext(io, :compact => true)
-    print(io, "ℤ[", b.lower, ":", b.upper, "]")
+Represents the n-fold Cartesian product of a set.
+"""
+struct CartesianPower{S,A} <: AbstractCartSetProd
+    _base::S
+    _axes::A
 end
 
-testvalue(b::BoundedInts) = min(b.lower, 0)
+@inline setcartpower(s::SetLike, dims) = CartesianPower(s, asaxes(dims))
 
-function Base.getindex(::typeof(ℤ), r::AbstractUnitRange)
-    BoundedInts(extrema(r)...)
+@inline pwr_base(s::CartesianPower) = s._base
+@inline pwr_axes(s::CartesianPower) = s._axes
+@inline pwr_size(s::CartesianPower) = axes2size(s.axes)
+
+componentsets(d::CartesianPower) = fill_with(d.parent, d.axes)
+
+function Base.in(x::AbstractArray, s::CartesianPower)
+    axes2size(s.axes) == size(x) ||
+        throw(ArgumentError("Size of CartesianPower and given point are incompatible."))
+    isempty(x) ? true : all(Base.Fix1(in, s.parent), x)::Bool
 end
 
+Base.isempty(s::CartesianPower) = isempty(s.parent) || size2length(axes2size(s.axes)) == 0
 
+function Base.union(s::CartesianPower, others::CartesianPower...)
+    axs = s.axes
 
-function Base.in(x::AbstractArray{T}, ::Simplex) where {T}
-    all(≥(zero(eltype(x))), x) || return false
-    return x ∈ zeroset(Simplex())
+    all(isequal(axs), map(x -> x.axes, others)) || throw(
+        ArgumentError("Cannot create union of CartesianPower sets with different axes."),
+    )
+
+    setcartpower(union(s.parent, map(x -> x.parent, others)...))
 end
 
 
+"""
+    struct CombinedSet <: ValueSet
+
+Represents a combination of two sets.
 
-struct Sphere <: CodimOne end
+User code should not create instances of `CombinedMeasure` directly, but should call
+[`combinesets(f_c, α, β)`](@ref) instead.
+"""
 
-function Base.in(x::AbstractArray{T}, ::Sphere) where {T}
-    return x ∈ zeroset(Sphere())
+struct CombinedSet{FC,MA<:SetLike,MB<:SetLike} <: ValueSet
+    f_c::FC
+    α::MA
+    β::MB
 end
diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl
index c61d0624..31398f32 100644
--- a/src/primitives/counting.jl
+++ b/src/primitives/counting.jl
@@ -40,3 +40,8 @@ insupport(μ::Counting{T}, x) where {T<:Type} = x isa μ.support
 massof(c::Counting, s::Set) = massof(CountingBase(), filter(insupport(c), s))
 
 massof(::CountingBase, s::Set) = length(s)
+
+# ToDo: Would this be correct?
+# @inline mdomain(::CountingBase) = IntegerValues()
+
+@inline mdomain(::Counting{DomainType}) where {DomainType} = DomainType()
diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl
index 3846eaf5..040d5bd2 100644
--- a/src/primitives/lebesgue.jl
+++ b/src/primitives/lebesgue.jl
@@ -75,23 +75,23 @@ massof(::Lebesgue{RealNumbers}, s::Interval) = width(s)
 # Example: 
 # julia> Lebesgue(𝕀)(0.2..5)
 # 0.8
-function massof(μ::Lebesgue{<:BoundedReals}, s::Interval)
-    a = μ.support.lower
-    b = μ.support.upper
+function massof(μ::Lebesgue{<:AbstractInterval}, s::Interval)
+    a, b = endpoinnts(μ.support)
     left = max(s.left, a)
     right = min(s.right, b)
     w = right - left
     max(w, zero(w))
 end
 
-function smf(μ::Lebesgue{<:BoundedReals}, x)
-    clamp(x, μ.support.lower, μ.support.upper)
+function smf(μ::Lebesgue{<:AbstractInterval}, x)
+    a, b = endpoinnts(μ.support)
+    clamp(x, a, b)
 end
 
-smf(::Lebesgue{RealNumbers}, x) = x
-smf(::Lebesgue{RealNumbers}) = identity
-invsmf(::Lebesgue{RealNumbers}, x) = x
-invsmf(::Lebesgue{RealNumbers}) = identity
+smf(::Lebesgue{<:RealNumbers}, x) = x
+smf(::Lebesgue{<:RealNumbers}) = identity
+invsmf(::Lebesgue{<:RealNumbers}, x) = x
+invsmf(::Lebesgue{<:RealNumbers}) = identity
 
 smf(::LebesgueBase, x) = x
 smf(::LebesgueBase) = identity

From bf2ad45bfd35a284f67a39ebd663fa7b12e04312 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 14:32:27 +0200
Subject: [PATCH 018/122] Complete domain redesign

---
 src/domains.jl             | 67 ++++++++++++++++++++++++++++++--------
 src/primitives/lebesgue.jl | 18 +++++-----
 src/standard/stduniform.jl |  2 +-
 3 files changed, 64 insertions(+), 23 deletions(-)

diff --git a/src/domains.jl b/src/domains.jl
index c69cc29e..7458b450 100644
--- a/src/domains.jl
+++ b/src/domains.jl
@@ -30,13 +30,13 @@ See also [`MeasureBase.SetLike`](@ref).
 abstract type ValueSet end
 
 """
-    const MeasureBase.SetLike = Union{MeaureBase.ValueSet, Base.AbstractSet, IntervalSets.Domain}
+    const MeasureBase.SetLike = Union{MeasureBase.ValueSet, Base.AbstractSet, IntervalSets.Domain}
 
 Any kind of (measurable) set.
 
 There needs to be an implicit sigma-algebra for subtypes of
 `MeasureBase.SetLike` to make them useable for measures. This can't easily be
-imposed via type constraints, though, to is is by-contract.
+imposed via type constraints, though, so it is by-contract.
 """
 const SetLike = Union{MeasureBase.ValueSet,Base.AbstractSet,IntervalSets.Domain}
 
@@ -142,7 +142,7 @@ maybe_in(@nospecialize(x), ::UnknownDomain) = true
 Base.isempty(::UnknownDomain) = false
 
 """
-    RealInterval() isa MeasureBase.ValueSet
+    RealValues() isa MeasureBase.ValueSet
 
 The real numbers.
 """
@@ -158,13 +158,17 @@ struct RealValues <: ValueSet end
 @inline Base.minimum(::RealValues) = static(-Inf)
 @inline Base.maximum(::RealValues) = static(Inf)
 
+testvalue(::Type{T}, ::RealValues) where {T} = zero(T)
+
 """
     const MeasureBase.ℝ = RealValues()
 
 The set of all real numbers, see [`MeasureBase.RealValues`](@ref).
 """
 const ℝ = RealValues()
+export ℝ
 
+Base.show(io::IO, ::RealValues) = print(io, "ℝ")
 Base.show(io::IO, ::MIME"text/plain", ::RealValues) = print(io, "MeasureBase.ℝ")
 
 """
@@ -179,6 +183,8 @@ struct IntegerValues <: ValueSet end
 
 @inline Base.union(s::IntegerValues, ::IntegerValues...) = s
 
+testvalue(::Type{T}, ::IntegerValues) where {T} = zero(T)
+
 # # This could get tricky with mixed-precision code. Probably needs some
 # # special AbstractInteger infinity type (but custom AbstractInteger types
 # # may cause a lot of method invalidations, which is why Static.StaticInteger
@@ -192,9 +198,45 @@ struct IntegerValues <: ValueSet end
 The set of all integers, see [`MeasureBase.IntegerValues`](@ref).
 """
 const ℤ = IntegerValues()
+export ℤ
 
+Base.show(io::IO, ::IntegerValues) = print(io, "ℤ")
 Base.show(io::IO, ::MIME"text/plain", ::IntegerValues) = print(io, "MeasureBase.ℤ")
 
+"""
+    struct MeasureBase.BoundedInts{L,U} <: MeasureBase.ValueSet
+
+The integers from `lower` to `upper` (bounds may be infinite).
+
+Constructors:
+
+```julia
+BoundedInts(lower, upper)
+ℤ[lower:upper]
+```
+"""
+struct BoundedInts{L,U} <: ValueSet
+    lower::L
+    upper::U
+end
+
+@inline Base.in(x, b::BoundedInts) = x ∈ ℤ && b.lower <= x <= b.upper
+
+Base.isempty(b::BoundedInts) = b.lower > b.upper
+
+Base.minimum(b::BoundedInts) = b.lower
+Base.maximum(b::BoundedInts) = b.upper
+
+function Base.show(io::IO, b::BoundedInts)
+    io = IOContext(io, :compact => true)
+    print(io, "ℤ[", b.lower, ":", b.upper, "]")
+end
+
+testvalue(b::BoundedInts) = convert(Int, clamp(0, dynamic(b.lower), dynamic(b.upper)))
+testvalue(::Type{T}, b::BoundedInts) where {T} = convert(T, testvalue(b))
+
+Base.getindex(::typeof(ℤ), r::AbstractUnitRange) = BoundedInts(extrema(r)...)
+
 """
     struct MeasureBase.AbstractCartSetProd <: ValueSet
 
@@ -240,7 +282,7 @@ function Base.in(
     isempty(x) && isempty(sets) ? true : all(in.(x, sets))::Bool
 end
 
-@inline Base.isempty(s::CartesianProduct) = all(!isempty, componentsets(s))
+@inline Base.isempty(s::CartesianProduct) = any(isempty, componentsets(s))
 
 @inline function Base.union(
     s::CartesianProduct{<:Tuple{Vararg{Any,N}}},
@@ -277,26 +319,26 @@ end
 
 @inline pwr_base(s::CartesianPower) = s._base
 @inline pwr_axes(s::CartesianPower) = s._axes
-@inline pwr_size(s::CartesianPower) = axes2size(s.axes)
+@inline pwr_size(s::CartesianPower) = axes2size(pwr_axes(s))
 
-componentsets(d::CartesianPower) = fill_with(d.parent, d.axes)
+componentsets(s::CartesianPower) = maybestatic_fill(pwr_base(s), pwr_axes(s))
 
 function Base.in(x::AbstractArray, s::CartesianPower)
-    axes2size(s.axes) == size(x) ||
+    pwr_size(s) == size(x) ||
         throw(ArgumentError("Size of CartesianPower and given point are incompatible."))
-    isempty(x) ? true : all(Base.Fix1(in, s.parent), x)::Bool
+    isempty(x) ? true : all(Base.Fix1(in, pwr_base(s)), x)::Bool
 end
 
-Base.isempty(s::CartesianPower) = isempty(s.parent) || size2length(axes2size(s.axes)) == 0
+Base.isempty(s::CartesianPower) = isempty(pwr_base(s)) || size2length(pwr_size(s)) == 0
 
 function Base.union(s::CartesianPower, others::CartesianPower...)
-    axs = s.axes
+    axs = pwr_axes(s)
 
-    all(isequal(axs), map(x -> x.axes, others)) || throw(
+    all(isequal(axs), map(pwr_axes, others)) || throw(
         ArgumentError("Cannot create union of CartesianPower sets with different axes."),
     )
 
-    setcartpower(union(s.parent, map(x -> x.parent, others)...))
+    setcartpower(union(pwr_base(s), map(pwr_base, others)...), axs)
 end
 
 
@@ -308,7 +350,6 @@ Represents a combination of two sets.
 User code should not create instances of `CombinedMeasure` directly, but should call
 [`combinesets(f_c, α, β)`](@ref) instead.
 """
-
 struct CombinedSet{FC,MA<:SetLike,MB<:SetLike} <: ValueSet
     f_c::FC
     α::MA
diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl
index 040d5bd2..3d92a2ed 100644
--- a/src/primitives/lebesgue.jl
+++ b/src/primitives/lebesgue.jl
@@ -51,7 +51,7 @@ Lebesgue() = Lebesgue(ℝ)
 testvalue(::Type{T}, d::Lebesgue) where {T} = testvalue(T, d.support)::T
 
 proxy(d::Lebesgue) = restrict(in(d.support), LebesgueBase())
-proxy(::Lebesgue{MeasureBase.RealNumbers}) = LebesgueBase()
+proxy(::Lebesgue{MeasureBase.RealValues}) = LebesgueBase()
 
 @useproxy Lebesgue
 
@@ -61,7 +61,7 @@ Base.show(io::IO, d::Lebesgue) = print(io, "Lebesgue(", d.support, ")")
 
 insupport(μ::Lebesgue, x) = x ∈ μ.support
 
-insupport(::Lebesgue{RealNumbers}, ::Real) = true
+insupport(::Lebesgue{RealValues}, ::Real) = true
 
 @inline function logdensityof(μ::Lebesgue, x::Real)
     R = float(typeof(x))
@@ -70,13 +70,13 @@ end
 
 @inline logdensityof(μ::Lebesgue, x) = insupport(μ, x) ? 0.0 : -Inf
 
-massof(::Lebesgue{RealNumbers}, s::Interval) = width(s)
+massof(::Lebesgue{RealValues}, s::Interval) = width(s)
 
 # Example: 
 # julia> Lebesgue(𝕀)(0.2..5)
 # 0.8
 function massof(μ::Lebesgue{<:AbstractInterval}, s::Interval)
-    a, b = endpoinnts(μ.support)
+    a, b = endpoints(μ.support)
     left = max(s.left, a)
     right = min(s.right, b)
     w = right - left
@@ -84,14 +84,14 @@ function massof(μ::Lebesgue{<:AbstractInterval}, s::Interval)
 end
 
 function smf(μ::Lebesgue{<:AbstractInterval}, x)
-    a, b = endpoinnts(μ.support)
+    a, b = endpoints(μ.support)
     clamp(x, a, b)
 end
 
-smf(::Lebesgue{<:RealNumbers}, x) = x
-smf(::Lebesgue{<:RealNumbers}) = identity
-invsmf(::Lebesgue{<:RealNumbers}, x) = x
-invsmf(::Lebesgue{<:RealNumbers}) = identity
+smf(::Lebesgue{<:RealValues}, x) = x
+smf(::Lebesgue{<:RealValues}) = identity
+invsmf(::Lebesgue{<:RealValues}, x) = x
+invsmf(::Lebesgue{<:RealValues}) = identity
 
 smf(::LebesgueBase, x) = x
 smf(::LebesgueBase) = identity
diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl
index 7bbe15ed..e3702656 100644
--- a/src/standard/stduniform.jl
+++ b/src/standard/stduniform.jl
@@ -14,7 +14,7 @@ end
 
 Base.rand(rng::Random.AbstractRNG, ::Type{T}, ::StdUniform) where {T} = rand(rng, T)
 
-massof(::StdUniform, s::Interval) = massof(Lebesgue(𝕀), s::Interval)
+massof(::StdUniform, s::Interval) = massof(Lebesgue(0.0 .. 1.0), s)
 
 smf(::StdUniform, x) = clamp(x, zero(x), one(x))
 

From 7835f8492d2e2564eb8c27f4fc1864c06a66e0ba Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:10 +0200
Subject: [PATCH 019/122] Add mreshape

---
 src/MeasureBase.jl          |  1 +
 src/combinators/power.jl    |  3 ++
 src/combinators/reshape.jl  | 69 +++++++++++++++++++++++++++++++++++++
 test/combinators/reshape.jl |  7 ++++
 test/runtests.jl            |  1 +
 5 files changed, 81 insertions(+)
 create mode 100644 src/combinators/reshape.jl
 create mode 100644 test/combinators/reshape.jl

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index c4c0d955..4ab59601 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -184,6 +184,7 @@ include("primitives/trivial.jl")
 
 include("combinators/bind.jl")
 include("combinators/transformedmeasure.jl")
+include("combinators/reshape.jl")
 include("combinators/weighted.jl")
 include("combinators/superpose.jl")
 include("combinators/product.jl")
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 4d4760ba..6f065c3f 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -179,3 +179,6 @@ function logdensity_def(
 ) where {P<:PrimitiveMeasure,N}
     static(0.0)
 end
+
+
+@inline mspace_elsize(m::PowerMeasure) = axes2size(m.axes)
diff --git a/src/combinators/reshape.jl b/src/combinators/reshape.jl
new file mode 100644
index 00000000..dddae55b
--- /dev/null
+++ b/src/combinators/reshape.jl
@@ -0,0 +1,69 @@
+# ToDo: Support static resizes for static arrays
+
+"""
+    struct MeasureBase.Reshape <: Function
+
+Represents a function that reshapes an array.
+
+Supports `InverseFunctions.inverse` and
+`ChangesOfVariables.with_logabsdet_jacobian`.
+
+Constructor:
+
+```julia
+Reshape(output_size::Dims, input_size::Dims)
+```
+"""
+struct Reshape{M<:SizeLike,N<:SizeLike} <: Function
+    output_size::M
+    input_size::N
+
+    Reshape{M,N}(out_sz::M, in_sz::N) where {M<:SizeLike,N<:SizeLike} =
+        new{M,N}(out_sz, in_sz)
+end
+
+function Reshape(output_size::SizeLike, input_size::SizeLike)
+    out_sz = canonical_size(output_size)
+    in_sz = canonical_size(input_size)
+    return Reshape{typeof(out_sz), typeof(in_sz)}(out_sz, in_sz)
+end
+
+_throw_reshape_mismatch(sz, sz_x) = throw(DimensionMismatch("Reshape input size is $sz but got input of size $sz_x"))
+
+function (f::Reshape)(x::AbstractArray)
+    sz_x = maybestatic_size(x)
+    f.input_size == sz_x || _throw_reshape_mismatch(f.input_size, sz_x)
+    return reshape(x, f.output_size)
+end
+
+InverseFunctions.inverse(f::Reshape{M,N}) where {M,N} = Reshape{N,M}(f.input_size, f.output_size)
+
+function ChangesOfVariables.with_logabsdet_jacobian(f::Reshape, x::AbstractArray)
+    return f(x), zero(real_numtype(typeof(x)))
+end
+
+
+"""
+    mreshape(m::AbstractMeasure, sz::Vararg{N,IntegerLike}) where N
+    mreshape(m::AbstractMeasure, sz::NTuple{N,IntegerLike}) where N
+
+Reshape a measure `m` over an array-valued space, returning a measure over
+a space of arrays with shape `sz`.
+"""
+function mreshape end
+
+mreshape(m::AbstractMeasure, sz::IntegerLike...) = mreshape(m, sz)
+mreshape(m::AbstractMeasure, sz::SizeLike) = pushfwd(Reshape(sz, mspace_elsize(m)), m)
+
+
+"""
+    MeasureBase.mspace_elsize(m::AbstractMeasure)::MeasureBase.SizeLike
+
+Return the size of the elements of the measurable space of `m`.
+
+Defaults to the size of a test value of `m`, may be specialized for
+measure types where this is inefficient.
+"""
+function mspace_elsize end
+
+mspace_elsize(m::AbstractMeasure) = maybestatic_size(testvalue(m))
diff --git a/test/combinators/reshape.jl b/test/combinators/reshape.jl
new file mode 100644
index 00000000..c6624582
--- /dev/null
+++ b/test/combinators/reshape.jl
@@ -0,0 +1,7 @@
+using Test
+
+using MeasureBase
+
+@testset "reshape" begin
+
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index c2f63c4e..dfd8b93a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -21,6 +21,7 @@ include("smf.jl")
 
 include("combinators/weighted.jl")
 include("combinators/transformedmeasure.jl")
+include("combinators/reshape.jl")
 include("combinators/implicitlymapped.jl")
 
 include("test_docs.jl")

From 78fc01ec4f95982ca72221368e50478027a71207 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 00:49:49 +0200
Subject: [PATCH 020/122] Require SpecialFunctions 2.1.4

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 7dea2fc6..e729f511 100644
--- a/Project.toml
+++ b/Project.toml
@@ -71,7 +71,7 @@ PrettyPrinting = "0.3, 0.4"
 PropertyFunctions = "0.2.2"
 Random = "1"
 Reexport = "1"
-SpecialFunctions = "2"
+SpecialFunctions = "2.1.4"
 Static = "0.8, 1"
 StaticArrays = "1.5"
 StaticThings = "0.2"

From 9853449b5118b0d05f83341f46a48344c4480010 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 00:50:31 +0200
Subject: [PATCH 021/122] Fix _default_checked_arg

---
 src/getdof.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/getdof.jl b/src/getdof.jl
index dbce2202..16ae7cc6 100644
--- a/src/getdof.jl
+++ b/src/getdof.jl
@@ -66,7 +66,7 @@ function checked_arg end
 
 # Prevent infinite recursion:
 @propagate_inbounds function _default_checked_arg(::Type{MU}, ::MU, ::T) where {MU,T}
-    NoArgCheck{MU,T}
+    NoArgCheck{MU,T}()
 end
 @propagate_inbounds function _default_checked_arg(::Type{MU}, mu_base, x) where {MU}
     checked_arg(mu_base, x)

From b9c53415b1385bfc892944f9cb4d90e15f6dbd34 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 00:52:05 +0200
Subject: [PATCH 022/122] Fix checked_arg for ProductMeasure

---
 src/combinators/product.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 0290419d..9135dc2b 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -244,6 +244,10 @@ function checked_arg(μ::ProductMeasure{<:NTuple{N,Any}}, x::NTuple{N,Any}) wher
     map(checked_arg, marginals(μ), x)
 end
 
+function checked_arg(μ::ProductMeasure{<:AbstractArray}, x::AbstractArray)
+    map(checked_arg, marginals(μ), x)
+end
+
 function checked_arg(
     μ::ProductMeasure{<:NamedTuple{names}},
     x::NamedTuple{names},

From 87bde7e4bc81cfa1e804e1a45b093ecbbb23c4bb Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 00:52:48 +0200
Subject: [PATCH 023/122] More transport_def methods for PowerMeasure and
 ProductMeasure

---
 src/standard/stdmeasure.jl | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index e7244fac..a9c09b5c 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -121,3 +121,28 @@ function transport_def(
 ) where {MU<:StdMeasure,names}
     NamedTuple{names}(_tuple_transport_def(values(marginals(ν)), μ, x))
 end
+
+function transport_def(
+    ν::PowerMeasure{NU},
+    μ::ProductMeasure{<:AbstractArray},
+    x,
+) where {NU<:StdMeasure}
+    reshape(vcat(map(_TransportToStd{NU}(), marginals(μ), x)...), ν.axes)
+end
+
+function _marginal_viewranges(μs::AbstractArray, startidx::IntegerLike)
+    ns = map(m -> dynamic(getdof(m)), μs)
+    offs = cumsum(vcat(dynamic(startidx), ns[begin:(end-1)]))
+    map((o, n) -> o:(o+n-1), offs, ns)
+end
+
+function transport_def(
+    ν::ProductMeasure{<:AbstractArray},
+    μ::PowerMeasure{MU},
+    x::AbstractArray{<:Real},
+) where {MU<:StdMeasure}
+    νs = marginals(ν)
+    vrs = _marginal_viewranges(νs, firstindex(x))
+    xs = map(r -> view(x, r), vrs)
+    map(_TransportFromStd{MU}, νs, xs)
+end

From 43d355ffcbfb2edd92e0d50b5ad305bc3735d15a Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 08:37:30 +0200
Subject: [PATCH 024/122] Add Mooncake extension with AD rules

Created by generative AI.
---
 Project.toml                  |  3 ++
 ext/MeasureBaseMooncakeExt.jl | 29 ++++++++++++++++++
 test/Project.toml             |  1 +
 test/runtests.jl              |  2 ++
 test/test_mooncake.jl         | 56 +++++++++++++++++++++++++++++++++++
 5 files changed, 91 insertions(+)
 create mode 100644 ext/MeasureBaseMooncakeExt.jl
 create mode 100644 test/test_mooncake.jl

diff --git a/Project.toml b/Project.toml
index e729f511..486afa73 100644
--- a/Project.toml
+++ b/Project.toml
@@ -37,6 +37,7 @@ Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 
 [extensions]
 MeasureBaseChainRulesCoreExt = "ChainRulesCore"
@@ -44,6 +45,7 @@ MeasureBaseDistributionsExt = "Distributions"
 MeasureBaseDistributionsChainRulesCoreExt = ["Distributions", "ChainRulesCore"]
 MeasureBaseDistributionsForwardDiffExt = ["Distributions", "ForwardDiff"]
 MeasureBaseForwardDiffExt = "ForwardDiff"
+MeasureBaseMooncakeExt = "Mooncake"
 
 [compat]
 ChainRulesCore = "1"
@@ -66,6 +68,7 @@ LinearAlgebra = "1"
 LogExpFunctions = "0.3, 1"
 LogarithmicNumbers = "1"
 MappedArrays = "0.4"
+Mooncake = "0.5.34"
 NaNMath = "0.3, 1"
 PrettyPrinting = "0.3, 0.4"
 PropertyFunctions = "0.2.2"
diff --git a/ext/MeasureBaseMooncakeExt.jl b/ext/MeasureBaseMooncakeExt.jl
new file mode 100644
index 00000000..e401cb7d
--- /dev/null
+++ b/ext/MeasureBaseMooncakeExt.jl
@@ -0,0 +1,29 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseMooncakeExt
+
+using MeasureBase
+import Mooncake
+using Mooncake: @zero_derivative, MinimalCtx
+
+using MeasureBase: isneginf, isposinf, _adignore_call
+using MeasureBase: check_dof, require_insupport, _origin_depth
+using MeasureBase: logdensityof_rt
+
+# Unlike Zygote, Mooncake differentiates the collection utilities
+# (`_pushfront`, etc., mutating code in general), `checked_arg` and
+# `_checksupport` natively, so only the non-differentiable functions
+# need rules:
+
+@zero_derivative MinimalCtx Tuple{typeof(isneginf),Any}
+@zero_derivative MinimalCtx Tuple{typeof(isposinf),Any}
+
+@zero_derivative MinimalCtx Tuple{typeof(_adignore_call),Any}
+
+@zero_derivative MinimalCtx Tuple{typeof(require_insupport),Any,Any}
+@zero_derivative MinimalCtx Tuple{typeof(_origin_depth),Any}
+@zero_derivative MinimalCtx Tuple{typeof(check_dof),Any,Any}
+
+@zero_derivative MinimalCtx Tuple{typeof(logdensityof_rt),Any,Any}
+
+end # module MeasureBaseMooncakeExt
diff --git a/test/Project.toml b/test/Project.toml
index 376c1b05..fec229fd 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -12,6 +12,7 @@ IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 PropertyFunctions = "09e99361-2bb8-48a2-a80f-de58f0739eb4"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
diff --git a/test/runtests.jl b/test/runtests.jl
index dfd8b93a..ff5c3140 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -19,6 +19,8 @@ include("getdof.jl")
 include("transport.jl")
 include("smf.jl")
 
+include("test_mooncake.jl")
+
 include("combinators/weighted.jl")
 include("combinators/transformedmeasure.jl")
 include("combinators/reshape.jl")
diff --git a/test/test_mooncake.jl b/test/test_mooncake.jl
new file mode 100644
index 00000000..95ce84b5
--- /dev/null
+++ b/test/test_mooncake.jl
@@ -0,0 +1,56 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using Random
+import Mooncake
+import ForwardDiff
+
+using MeasureBase
+using MeasureBase: transport_to
+using MeasureBase: isneginf, isposinf, _adignore_call
+using MeasureBase: check_dof, require_insupport, _origin_depth
+using MeasureBase: logdensityof_rt
+
+_mooncake_gradient(f, x) = Mooncake.value_and_gradient!!(
+    Mooncake.prepare_gradient_cache(f, x), f, x
+)[2][2]
+
+@testset "Mooncake AD rules" begin
+    @test Base.get_extension(MeasureBase, :MeasureBaseMooncakeExt) isa Module
+
+    @testset "zero-derivative primitives" begin
+        rng = Random.Xoshiro(789990641)
+        Mooncake.TestUtils.test_rule(rng, isneginf, 0.5; is_primitive = true)
+        Mooncake.TestUtils.test_rule(rng, isposinf, 0.5; is_primitive = true)
+        Mooncake.TestUtils.test_rule(rng, _adignore_call, () -> 42.0; is_primitive = true)
+        Mooncake.TestUtils.test_rule(rng, check_dof, StdNormal(), StdUniform(); is_primitive = true)
+        Mooncake.TestUtils.test_rule(rng, require_insupport, StdNormal(), 0.5; is_primitive = true)
+        Mooncake.TestUtils.test_rule(rng, _origin_depth, StdNormal(); is_primitive = true)
+        Mooncake.TestUtils.test_rule(rng, logdensityof_rt, StdNormal(), 0.5; is_primitive = true)
+    end
+
+    @testset "@_adignore is ignored" begin
+        f_adignore(x) = (MeasureBase.@_adignore x^3; x^2)
+        @test _mooncake_gradient(f_adignore, 3.0) ≈ 6.0
+    end
+
+    @testset "logdensityof gradients" begin
+        x = [0.1, -0.2, 0.3]
+        f_ld = x -> logdensityof(StdNormal()^3, x)
+        @test _mooncake_gradient(f_ld, x) ≈ ForwardDiff.gradient(f_ld, x)
+
+        f_ldu = x -> logdensityof(StdExponential()^3, x)
+        @test _mooncake_gradient(f_ldu, abs.(x)) ≈ ForwardDiff.gradient(f_ldu, abs.(x))
+    end
+
+    @testset "transport gradients" begin
+        x = [0.1, -0.2, 0.3]
+        f_t = x -> sum(transport_to(StdUniform()^3, StdNormal()^3)(x))
+        @test _mooncake_gradient(f_t, x) ≈ ForwardDiff.gradient(f_t, x)
+
+        u = [0.3, 0.5, 0.7]
+        f_ti = u -> sum(transport_to(StdNormal()^3, StdUniform()^3)(u))
+        @test _mooncake_gradient(f_ti, u) ≈ ForwardDiff.gradient(f_ti, u)
+    end
+end

From ba5e85f92c6411ba68de7e3ca6123d8e2b7abed0 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sun, 5 Jul 2026 12:40:10 +0200
Subject: [PATCH 025/122] Add Distributions extension

Assisted by generative AI.
---
 Project.toml                                  |  20 +-
 ext/MeasureBaseChainRulesCoreExt.jl           |  15 +-
 ...asureBaseDistributionsChainRulesCoreExt.jl |   9 +
 ext/MeasureBaseDistributionsExt.jl            |   8 -
 .../MeasureBaseDistributionsExt.jl            |  66 ++++++
 ext/MeasureBaseDistributionsExt/dirac.jl      |  14 ++
 ext/MeasureBaseDistributionsExt/dirichlet.jl  |  59 +++++
 .../dist_vartransform.jl                      |  16 ++
 .../distribution_measure.jl                   |  74 ++++++
 .../measure_interface.jl                      |  26 ++
 ext/MeasureBaseDistributionsExt/mixture.jl    |  24 ++
 ext/MeasureBaseDistributionsExt/product.jl    |  56 +++++
 ext/MeasureBaseDistributionsExt/reshaped.jl   |  27 +++
 .../standard_dist.jl                          | 222 ++++++++++++++++++
 .../standard_normal.jl                        |  77 ++++++
 .../standard_uniform.jl                       |  79 +++++++
 ext/MeasureBaseDistributionsExt/standardmv.jl |  36 +++
 ext/MeasureBaseDistributionsExt/univariate.jl | 142 +++++++++++
 ext/MeasureBaseDistributionsForwardDiffExt.jl |  29 ++-
 ...aseDistributionsForwardDiffPullbacksExt.jl |  25 ++
 ext/MeasureBaseDistributionsMooncakeExt.jl    |  21 ++
 ext/MeasureBaseForwardDiffExt.jl              |   4 +
 ext/MeasureBaseForwardDiffPullbacksExt.jl     |  10 +
 src/combinators/superpose.jl                  | 102 ++++----
 src/utils.jl                                  |  51 ++++
 test/Project.toml                             |   7 +
 test/distributions/getjacobian.jl             |  34 +++
 test/distributions/test_autodiff_utils.jl     |  18 ++
 test/distributions/test_conversions.jl        | 113 +++++++++
 .../test_distribution_measure.jl              |  53 +++++
 test/distributions/test_distributions.jl      |  24 ++
 test/distributions/test_measure_interface.jl  |  43 ++++
 test/distributions/test_mooncake.jl           |  68 ++++++
 test/distributions/test_standard_dist.jl      | 128 ++++++++++
 test/distributions/test_standard_normal.jl    | 129 ++++++++++
 test/distributions/test_standard_uniform.jl   | 118 ++++++++++
 test/distributions/test_transport.jl          | 194 +++++++++++++++
 test/runtests.jl                              |   3 +-
 test/test_aqua.jl                             |   7 +-
 39 files changed, 2080 insertions(+), 71 deletions(-)
 delete mode 100644 ext/MeasureBaseDistributionsExt.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/dirac.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/dirichlet.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/dist_vartransform.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/distribution_measure.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/measure_interface.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/mixture.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/product.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/reshaped.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/standard_dist.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/standard_normal.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/standard_uniform.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/standardmv.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/univariate.jl
 create mode 100644 ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl
 create mode 100644 ext/MeasureBaseDistributionsMooncakeExt.jl
 create mode 100644 ext/MeasureBaseForwardDiffPullbacksExt.jl
 create mode 100644 test/distributions/getjacobian.jl
 create mode 100644 test/distributions/test_autodiff_utils.jl
 create mode 100644 test/distributions/test_conversions.jl
 create mode 100644 test/distributions/test_distribution_measure.jl
 create mode 100644 test/distributions/test_distributions.jl
 create mode 100644 test/distributions/test_measure_interface.jl
 create mode 100644 test/distributions/test_mooncake.jl
 create mode 100644 test/distributions/test_standard_dist.jl
 create mode 100644 test/distributions/test_standard_normal.jl
 create mode 100644 test/distributions/test_standard_uniform.jl
 create mode 100644 test/distributions/test_transport.jl

diff --git a/Project.toml b/Project.toml
index 486afa73..e2bfee4a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,6 +4,8 @@ version = "0.14.12"
 authors = ["Chad Scherrer <chad.scherrer@gmail.com>", "Oliver Schulz <oschulz@mpp.mpg.de>", "contributors"]
 
 [deps]
+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"
@@ -37,17 +39,26 @@ Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 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"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 
 [extensions]
 MeasureBaseChainRulesCoreExt = "ChainRulesCore"
-MeasureBaseDistributionsExt = "Distributions"
+MeasureBaseDistributionsExt = ["Distributions", "StatsBase", "StatsFuns", "PDMats"]
 MeasureBaseDistributionsChainRulesCoreExt = ["Distributions", "ChainRulesCore"]
 MeasureBaseDistributionsForwardDiffExt = ["Distributions", "ForwardDiff"]
+MeasureBaseDistributionsForwardDiffPullbacksExt = ["Distributions", "ForwardDiffPullbacks", "ChainRulesCore"]
+MeasureBaseDistributionsMooncakeExt = ["Distributions", "Mooncake"]
 MeasureBaseForwardDiffExt = "ForwardDiff"
+MeasureBaseForwardDiffPullbacksExt = "ForwardDiffPullbacks"
 MeasureBaseMooncakeExt = "Mooncake"
 
 [compat]
+ArgCheck = "1, 2"
+ArraysOfArrays = "0.6"
 ChainRulesCore = "1"
 ChangesOfVariables = "0.1.3"
 Compat = "3.35, 4"
@@ -55,9 +66,9 @@ ConstantRNGs = "0.1.1"
 ConstructionBase = "1.3"
 DensityInterface = "0.4"
 Distributions = "0.25.1"
-Distributions = "0.25.111"
 FillArrays = "0.12, 0.13, 1"
-ForwardDiff = "0.8, 0.9, 0.10"
+ForwardDiff = "0.10, 1"
+ForwardDiffPullbacks = "0.2"
 FunctionChains = "0.2.3"
 HeterogeneousComputing = "0.2.3"
 IfElse = "0.1"
@@ -70,6 +81,7 @@ LogarithmicNumbers = "1"
 MappedArrays = "0.4"
 Mooncake = "0.5.34"
 NaNMath = "0.3, 1"
+PDMats = "0.11"
 PrettyPrinting = "0.3, 0.4"
 PropertyFunctions = "0.2.2"
 Random = "1"
@@ -79,6 +91,8 @@ Static = "0.8, 1"
 StaticArrays = "1.5"
 StaticThings = "0.2"
 Statistics = "1"
+StatsBase = "0.33, 0.34"
+StatsFuns = "0.9, 1, 2"
 Test = "1"
 Tricks = "0.1"
 julia = "1.10"
diff --git a/ext/MeasureBaseChainRulesCoreExt.jl b/ext/MeasureBaseChainRulesCoreExt.jl
index 0384a04b..25019da1 100644
--- a/ext/MeasureBaseChainRulesCoreExt.jl
+++ b/ext/MeasureBaseChainRulesCoreExt.jl
@@ -11,14 +11,25 @@ import ChainRulesCore
 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: _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log
+using MeasureBase: _pushfront, _pushback, _rev_cumsum, _exp_cumsum_log
 
 function ChainRulesCore.rrule(::typeof(_pushfront), v::AbstractVector, x)
     result = _pushfront(v, x)
diff --git a/ext/MeasureBaseDistributionsChainRulesCoreExt.jl b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl
index 4dd3f4ff..63cc3e93 100644
--- a/ext/MeasureBaseDistributionsChainRulesCoreExt.jl
+++ b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl
@@ -5,5 +5,14 @@ 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())
+function ChainRulesCore.rrule(::typeof(_dist_params_numtype), d::Distribution)
+    _dist_params_numtype(d), _dist_params_numtype_pullback
+end
 
 end # module MeasureBaseDistributionsChainRulesCoreExt
diff --git a/ext/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt.jl
deleted file mode 100644
index beb47821..00000000
--- a/ext/MeasureBaseDistributionsExt.jl
+++ /dev/null
@@ -1,8 +0,0 @@
-# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
-
-module MeasureBaseDistributionsExt
-
-using MeasureBase
-import Distributions
-
-end # module MeasureBaseDistributionsExt
diff --git a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
new file mode 100644
index 00000000..177ec306
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
@@ -0,0 +1,66 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseDistributionsExt
+
+using LinearAlgebra: Diagonal, diag, dot, cholesky
+
+import Random
+using Random: AbstractRNG, rand!
+
+import DensityInterface
+using DensityInterface: logdensityof, densityof
+
+import MeasureBase
+using MeasureBase: AbstractMeasure, AsMeasure, asmeasure
+using MeasureBase: Lebesgue, Counting, ℝ
+using MeasureBase: StdMeasure, StdUniform, StdExponential, StdLogistic, StdNormal
+using MeasureBase: PowerMeasure, WeightedMeasure, SuperpositionMeasure, PushforwardMeasure
+using MeasureBase: basemeasure, rootmeasure, testvalue, productmeasure, pushfwd, superpose
+using MeasureBase: getdof, checked_arg, massof
+using MeasureBase: transport_to, transport_def, transport_origin, from_origin, to_origin
+using MeasureBase: NoTransportOrigin, NoTransport
+using MeasureBase: Reshape
+using MeasureBase: convert_realtype, firsttype, _fwddiff, @_adignore
+import MeasureBase:
+    _dist_params_numtype, _trafo_cdf_impl, _trafo_quantile_impl, _trafo_quantile_impl_generic
+using MeasureBase: _pushfront, _pushback, _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log
+
+import Distributions
+using Distributions: Distribution, VariateForm, ValueSupport, ContinuousDistribution
+using Distributions: Univariate, Multivariate, ArrayLikeVariate, Continuous, Discrete
+using Distributions: Uniform, Exponential, Logistic, Normal
+using Distributions: MvNormal, AbstractMvNormal, Beta, Dirichlet
+using Distributions: ReshapedDistribution, AbstractMixtureModel
+
+import Statistics
+import StatsBase
+import StatsFuns
+import PDMats
+
+using IrrationalConstants: log2π, invsqrt2π
+
+using HeterogeneousComputing: real_numtype
+
+using Static: True, False, StaticInt, static, dynamic
+using StaticThings: asnonstatic
+using FillArrays: Fill, Ones, Zeros
+
+using ArgCheck: @argcheck
+
+using ArraysOfArrays: ArrayOfSimilarArrays, flatview
+
+include("measure_interface.jl")
+include("standard_dist.jl")
+include("standard_uniform.jl")
+include("standard_normal.jl")
+include("distribution_measure.jl")
+include("dist_vartransform.jl")
+include("univariate.jl")
+include("standardmv.jl")
+include("product.jl")
+include("reshaped.jl")
+include("mixture.jl")
+include("dirichlet.jl")
+include("dirac.jl")
+
+end # module MeasureBaseDistributionsExt
diff --git a/ext/MeasureBaseDistributionsExt/dirac.jl b/ext/MeasureBaseDistributionsExt/dirac.jl
new file mode 100644
index 00000000..8580df8c
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/dirac.jl
@@ -0,0 +1,14 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+MeasureBase.AbstractMeasure(obj::Distributions.Dirac) = MeasureBase.Dirac(obj.value)
+
+function AsMeasure{D}(::D) where {D<:Distributions.Dirac}
+    throw(ArgumentError("Don't wrap Distributions.Dirac into MeasureBase.AsMeasure, use asmeasure to convert instead."))
+end
+
+
+Distributions.Distribution(m::MeasureBase.Dirac{<:Real}) = Distributions.Dirac(m.x)
+
+function Distributions.Distribution(@nospecialize(m::MeasureBase.Dirac{T})) where T
+    throw(ArgumentError("Can only convert MeasureBase.Dirac{<:Real} to Distributions.Dirac, but not MeasureBase.Dirac{<:$(nameof(T))}"))
+end
diff --git a/ext/MeasureBaseDistributionsExt/dirichlet.jl b/ext/MeasureBaseDistributionsExt/dirichlet.jl
new file mode 100644
index 00000000..c60eeecd
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/dirichlet.jl
@@ -0,0 +1,59 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+const DirichletMeasure = AsMeasure{<:Dirichlet}
+
+MeasureBase.getdof(d::Dirichlet) = length(d) - 1
+MeasureBase.getdof(m::DirichletMeasure) = getdof(m.obj)
+
+MeasureBase.transport_origin(d::Dirichlet) = StdUniform()^getdof(d)
+
+
+
+function _dirichlet_beta_trafo(α::Real, β::Real, x::Real)
+    R = float(promote_type(typeof(α), typeof(β), typeof(x)))
+    convert(R, transport_def(Beta(α, β), StdUniform(), x))::R
+end
+
+_a_times_one_minus_b(a::Real, b::Real) = a * (1 - b)
+
+function MeasureBase.from_origin(ν::Dirichlet, x)
+    # See M. J. Betancourt, "Cruising The Simplex: Hamiltonian Monte Carlo and the Dirichlet Distribution",
+    # https://arxiv.org/abs/1010.3436
+
+    @_adignore @argcheck length(ν) == length(x) + 1
+
+    αs = _dropfront(_rev_cumsum(ν.alpha))
+    βs = _dropback(ν.alpha)
+    beta_v = _fwddiff(_dirichlet_beta_trafo).(αs, βs, x)
+    beta_v_cp = _exp_cumsum_log(_pushfront(beta_v, 1))
+    beta_v_ext = _pushback(beta_v, 0)
+    _fwddiff(_a_times_one_minus_b).(beta_v_cp, beta_v_ext)
+end
+
+
+function _inv_dirichlet_beta_trafo(α::Real, β::Real, beta_v::Real)
+    R = float(promote_type(typeof(α), typeof(β), typeof(beta_v)))
+    convert(R, transport_def(StdUniform(), Beta(α, β), beta_v))::R
+end
+
+# ToDo: Find efficient pullback for this:
+function _dirichlet_variate_to_beta_v(y::AbstractVector{<:Real})
+    beta_v = similar(y, length(eachindex(y)) - 1)
+    @assert firstindex(beta_v) == firstindex(y)
+    @assert lastindex(beta_v) == lastindex(y) - 1
+    T = eltype(y)
+    sum_log_beta_v::T = 0
+    @inbounds for i in eachindex(beta_v)
+        beta_v[i] = 1 - y[i] / exp(sum_log_beta_v)
+        sum_log_beta_v += log(beta_v[i])
+    end
+    return beta_v
+end
+
+function MeasureBase.to_origin(ν::Dirichlet, y)
+    @_adignore @argcheck length(ν) == length(y)
+    αs = _dropfront(_rev_cumsum(ν.alpha))
+    βs = _dropback(ν.alpha)
+    beta_v = _dirichlet_variate_to_beta_v(y)
+    _fwddiff(_inv_dirichlet_beta_trafo).(αs, βs, beta_v)
+end
diff --git a/ext/MeasureBaseDistributionsExt/dist_vartransform.jl b/ext/MeasureBaseDistributionsExt/dist_vartransform.jl
new file mode 100644
index 00000000..ceedabe9
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/dist_vartransform.jl
@@ -0,0 +1,16 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+const _AnyStdUniform = Union{StandardUniform,Uniform}
+const _AnyStdNormal = Union{StandardNormal,Normal}
+
+const _AnyStdDistribution = Union{_AnyStdUniform,_AnyStdNormal}
+
+_std_dist(::Type{<:_AnyStdUniform}) = StandardUniform
+_std_dist(::Type{<:_AnyStdNormal}) = StandardNormal
+
+_std_dist(::Type{D}, ::StaticInt{1}) where {D<:_AnyStdDistribution} = D()
+_std_dist(::Type{D}, dof) where {D<:_AnyStdDistribution} = D(dynamic(dof))
+_std_dist_for(::Type{D}, μ::Any) where {D<:_AnyStdDistribution} = _std_dist(_std_dist(D), getdof(μ))
+
+MeasureBase.transport_to(::Type{NU}, μ) where {NU<:_AnyStdDistribution} = transport_to(_std_dist_for(NU, μ), μ)
+MeasureBase.transport_to(ν, ::Type{MU}) where {MU<:_AnyStdDistribution} = transport_to(ν, _std_dist_for(MU, ν))
diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
new file mode 100644
index 00000000..bcbfd558
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -0,0 +1,74 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+
+const DistributionMeasure{F<:VariateForm,S<:ValueSupport,D<:Distribution{F,S}} = AsMeasure{D}
+
+@inline MeasureBase.AbstractMeasure(obj::Distribution) = AsMeasure{typeof(obj)}(obj)
+@inline Base.convert(::Type{AbstractMeasure}, obj::Distribution) = AbstractMeasure(obj)
+
+@inline Distributions.Distribution(m::DistributionMeasure) = m.obj
+@inline Distributions.Distribution{F}(m::DistributionMeasure{F}) where {F<:VariateForm} = Distribution(m)
+@inline Distributions.Distribution{F,S}(m::DistributionMeasure{F,S}) where {F<:VariateForm,S<:ValueSupport} = Distribution(m)
+
+@inline Base.convert(::Type{Distribution}, m::DistributionMeasure) = Distribution(m)
+@inline Base.convert(::Type{Distribution{F}}, m::DistributionMeasure{F}) where {F<:VariateForm} = Distribution(m)
+@inline Base.convert(::Type{Distribution{F,S}}, m::DistributionMeasure{F,S}) where {F<:VariateForm,S<:ValueSupport} = Distribution(m)
+
+
+Base.rand(rng::AbstractRNG, ::Type{T}, m::DistributionMeasure) where {T<:Real} = convert_realtype(T, rand(m.obj))
+
+function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution{<:ArrayLikeVariate{0}}, sz::Dims) where {T<:Real}
+    convert_realtype(T, reshape(rand(rng, d, prod(sz)), sz...))
+end
+
+function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution{<:ArrayLikeVariate{1}}, sz::Dims) where {T<:Real}
+    convert_realtype(T, reshape(rand(rng, d, prod(sz)), size(d)..., sz...))
+end
+
+function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::ReshapedDistribution{N,<:Any,<:Distribution{<:ArrayLikeVariate{1}}}, sz::Dims) where {T<:Real,N}
+    convert_realtype(T, reshape(rand(rng, d.dist, prod(sz)), d.dims..., sz...))
+end
+
+function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution, sz::Dims) where {T<:Real}
+    flatview(ArrayOfSimilarArrays(convert_realtype(T, rand(rng, d, sz))))
+end
+
+function Base.rand(rng::AbstractRNG, ::Type{T}, m::PowerMeasure{<:DistributionMeasure{<:ArrayLikeVariate{0}}, NTuple{N,Base.OneTo{Int}}}) where {T<:Real,N}
+    _flat_powrand(rng, T, m.parent.obj, map(length, m.axes))
+end
+
+function Base.rand(rng::AbstractRNG, ::Type{T}, m::PowerMeasure{<:DistributionMeasure{<:ArrayLikeVariate{M}}, NTuple{N,Base.OneTo{Int}}}) where {T<:Real,M,N}
+    flat_data = _flat_powrand(rng, T, m.parent.obj, map(length, m.axes))
+    ArrayOfSimilarArrays{T,M,N}(flat_data)
+end
+
+
+@inline DensityInterface.densityof(m::DistributionMeasure) = densityof(m.obj)
+@inline DensityInterface.logdensityof(m::DistributionMeasure) = logdensityof(m.obj)
+
+@inline MeasureBase.logdensity_def(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x)
+@inline MeasureBase.unsafe_logdensityof(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x)
+@inline MeasureBase.insupport(m::DistributionMeasure, x) = Distributions.insupport(m.obj, x)
+
+@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate{0},<:Continuous}) = Lebesgue()
+@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate,<:Continuous}) = Lebesgue()^size(m.obj)
+@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate{0},<:Discrete}) = Counting()
+@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate,<:Discrete}) = Counting()^size(m.obj)
+
+@inline MeasureBase.basemeasure(m::DistributionMeasure) = rootmeasure(m)
+
+@inline MeasureBase.massof(::DistributionMeasure) = static(1.0)
+
+@inline MeasureBase.mspace_elsize(m::DistributionMeasure{<:ArrayLikeVariate}) = size(m.obj)
+
+@inline MeasureBase.getdof(m::DistributionMeasure{<:ArrayLikeVariate{0}}) = 1
+
+# Delegate transport to the wrapped distribution:
+@inline MeasureBase.transport_origin(m::DistributionMeasure) = m.obj
+@inline MeasureBase.to_origin(::DistributionMeasure, y) = y
+@inline MeasureBase.from_origin(::DistributionMeasure, x) = x
+
+@inline MeasureBase.paramnames(m::DistributionMeasure) = propertynames(m.obj)
+@inline MeasureBase.params(m::DistributionMeasure) = NamedTuple{propertynames(m.obj)}(Distributions.params(m.obj))
+
+# @inline MeasureBase.testvalue(m::DistributionMeasure) = testvalue(basemeasure(d))
diff --git a/ext/MeasureBaseDistributionsExt/measure_interface.jl b/ext/MeasureBaseDistributionsExt/measure_interface.jl
new file mode 100644
index 00000000..6fed5d4a
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/measure_interface.jl
@@ -0,0 +1,26 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+@inline MeasureBase.logdensity_def(d::Distribution, x) = DensityInterface.logdensityof(d, x)
+@inline MeasureBase.unsafe_logdensityof(d::Distribution, x) = DensityInterface.logdensityof(d, x)
+
+@inline MeasureBase.insupport(d::Distribution, x) = Distributions.insupport(d, x)
+
+@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate{0},<:Continuous}) = Lebesgue()
+@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate,<:Continuous}) = Lebesgue()^size(d)
+@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate{0},<:Discrete}) = Counting()
+@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate,<:Discrete}) = Counting()^size(d)
+
+@inline MeasureBase.paramnames(d::Distribution) = propertynames(d)
+@inline MeasureBase.params(d::Distribution) = NamedTuple{propertynames(d)}(Distributions.params(d))
+
+@inline MeasureBase.testvalue(d::Distribution) = testvalue(basemeasure(d))
+@inline MeasureBase.testvalue(::Type{T}, d::Distribution) where {T} = testvalue(T, basemeasure(d))
+
+
+@inline MeasureBase.basemeasure(d::Distributions.Poisson) =
+    Counting(MeasureBase.BoundedInts(static(0), static(Inf)))
+@inline MeasureBase.basemeasure(d::Distributions.Product{<:Any,<:Distributions.Poisson}) =
+    Counting(MeasureBase.BoundedInts(static(0), static(Inf)))^size(d)
+
+
+MeasureBase.∫(f, base::Distribution) = MeasureBase.∫(f, convert(AbstractMeasure, base))
diff --git a/ext/MeasureBaseDistributionsExt/mixture.jl b/ext/MeasureBaseDistributionsExt/mixture.jl
new file mode 100644
index 00000000..89d41acf
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/mixture.jl
@@ -0,0 +1,24 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+function MeasureBase.AbstractMeasure(d::Distributions.AbstractMixtureModel)
+    superpose(map((w, c) -> w * asmeasure(c), Distributions.probs(d), Distributions.components(d)))
+end
+
+function AsMeasure{D}(::D) where {D<:Distributions.AbstractMixtureModel}
+    throw(ArgumentError("Don't wrap Distributions.AbstractMixtureModel into MeasureBase.AsMeasure, use asmeasure to convert instead."))
+end
+
+
+const _MixtureMeasure = SuperpositionMeasure{
+    <:Union{Tuple{Vararg{WeightedMeasure}},AbstractVector{<:WeightedMeasure}},
+}
+
+_mixture_component(m::AsMeasure{<:Distribution}) = m.obj
+
+function Distributions.Distribution(m::_MixtureMeasure)
+    components = map(c -> _mixture_component(c.base), collect(values(m.components)))
+    prior = map(c -> exp(c.logweight), collect(values(m.components)))
+    Distributions.MixtureModel(components, prior)
+end
+
+Base.convert(::Type{Distribution}, m::_MixtureMeasure) = Distributions.Distribution(m)
diff --git a/ext/MeasureBaseDistributionsExt/product.jl b/ext/MeasureBaseDistributionsExt/product.jl
new file mode 100644
index 00000000..a050dc97
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/product.jl
@@ -0,0 +1,56 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+@static if isdefined(Distributions, :Product)
+    MeasureBase.AbstractMeasure(obj::Distributions.Product) = productmeasure(map(asmeasure, obj.v))
+
+    function AsMeasure{D}(::D) where {D<:Distributions.Product}
+        throw(ArgumentError("Don't wrap Distributions.Product into MeasureBase.AsMeasure, use asmeasure to convert instead."))
+    end
+end
+
+function Distributions.Distribution(
+    m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution{Univariate}}}},
+)
+    Distributions.product_distribution(map(x -> x.obj, MeasureBase.marginals(m)))
+end
+
+function Base.convert(
+    ::Type{Distribution},
+    m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution{Univariate}}}},
+)
+    Distributions.Distribution(m)
+end
+
+@static if isdefined(Distributions, :ProductDistribution)
+    MeasureBase.AbstractMeasure(obj::Distributions.ProductDistribution) = productmeasure(map(asmeasure, obj.dists))
+
+    function AsMeasure{D}(::D) where {D<:Distributions.ProductDistribution}
+        throw(ArgumentError("Don't wrap Distributions.ProductDistribution into MeasureBase.AsMeasure, use asmeasure to convert instead."))
+    end
+
+    function Distributions.Distribution(
+        m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution}}},
+    )
+        Distributions.product_distribution(map(x -> x.obj, MeasureBase.marginals(m)))
+    end
+
+    function Distributions.Distribution(
+        m::MeasureBase.ProductMeasure{<:Tuple{Vararg{AsMeasure{<:Distribution}}}},
+    )
+        Distributions.product_distribution(map(x -> x.obj, MeasureBase.marginals(m))...)
+    end
+
+    function Base.convert(
+        ::Type{Distribution},
+        m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution}}},
+    )
+        Distributions.Distribution(m)
+    end
+
+    function Base.convert(
+        ::Type{Distribution},
+        m::MeasureBase.ProductMeasure{<:Tuple{Vararg{AsMeasure{<:Distribution}}}},
+    )
+        Distributions.Distribution(m)
+    end
+end
diff --git a/ext/MeasureBaseDistributionsExt/reshaped.jl b/ext/MeasureBaseDistributionsExt/reshaped.jl
new file mode 100644
index 00000000..6efde609
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/reshaped.jl
@@ -0,0 +1,27 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+MeasureBase.getdof(μ::ReshapedDistribution) = MeasureBase.getdof(μ.dist)
+
+MeasureBase.transport_origin(μ::ReshapedDistribution) = μ.dist
+
+MeasureBase.to_origin(ν::ReshapedDistribution, y) = reshape(y, size(ν.dist))
+
+MeasureBase.from_origin(ν::ReshapedDistribution, x) = reshape(x, ν.dims)
+
+
+function MeasureBase.AbstractMeasure(d::Distributions.ReshapedDistribution)
+    orig_dist = d.dist
+    pushfwd(Reshape(size(d), size(orig_dist)), AbstractMeasure(orig_dist))
+end
+
+function AsMeasure{D}(::D) where {D<:Distributions.ReshapedDistribution}
+    throw(ArgumentError("Don't wrap Distributions.ReshapedDistribution into MeasureBase.AsMeasure, use asmeasure to convert instead."))
+end
+
+
+function Distributions.Distribution(m::PushforwardMeasure{<:Reshape})
+    reshape(Distributions.Distribution(m.origin), asnonstatic(m.f.output_size)...)
+end
+
+Base.convert(::Type{Distribution}, m::PushforwardMeasure{<:Reshape}) =
+    Distributions.Distribution(m)
diff --git a/ext/MeasureBaseDistributionsExt/standard_dist.jl b/ext/MeasureBaseDistributionsExt/standard_dist.jl
new file mode 100644
index 00000000..010b20a9
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/standard_dist.jl
@@ -0,0 +1,222 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+"""
+    struct StandardDist{D<:Distribution{Univariate,Continuous},N} <: Distributions.Distribution{ArrayLikeVariate{N},Continuous}
+
+Represents `D()` or a product distribution of `D()` in a dispatchable fashion.
+
+Constructor:
+```
+    StandardDist{Uniform}(size...)
+    StandardDist{Normal}(size...)
+```
+"""
+struct StandardDist{D<:Distribution{Univariate,Continuous},N,U<:Integer} <:
+       Distributions.Distribution{ArrayLikeVariate{N},Continuous}
+    _size::NTuple{N,U}
+end
+export StandardDist
+
+StandardDist{D}() where {D<:Distribution{Univariate,Continuous}} =
+    StandardDist{D,0,Int}(())
+StandardDist{D}(dims::Vararg{U,N}) where {D<:Distribution{Univariate,Continuous},N,U<:Integer} =
+    StandardDist{D,N,U}((dims...,))
+
+
+const StandardUnivariateDist{D<:Distribution{Univariate,Continuous},U<:Integer} = StandardDist{D,0,U}
+const StandardMultivariteDist{D<:Distribution{Univariate,Continuous},U<:Integer} = StandardDist{D,1,U}
+
+
+function Base.show(io::IO, d::StandardDist{D}) where {D}
+    print(io, nameof(typeof(d)), "{", D, "}")
+    show(io, d._size)
+end
+
+
+@inline MeasureBase.transport_def(::MU, μ::MU, x) where {MU<:StandardDist{<:Any,0}} = x
+
+for (A, B) in [
+    (Uniform, StdUniform),
+    (Exponential, StdExponential),
+    (Logistic, StdLogistic),
+    (Normal, StdNormal)
+]
+    @eval begin
+        @inline MeasureBase.transport_origin(d::StandardDist{$A,0}) = $B()
+        @inline MeasureBase.transport_origin(d::StandardDist{$A,N}) where {N} = $B()^size(d)
+
+        # StandardDist{$A} and $B are equivalent as measures, so convert
+        # instead of wrapping:
+        MeasureBase.AbstractMeasure(::StandardDist{$A,0}) = $B()
+        MeasureBase.AbstractMeasure(d::StandardDist{$A,N}) where {N} = $B()^size(d)
+
+        Distributions.Distribution(::$B) = StandardDist{$A}()
+        Base.convert(::Type{Distribution}, ::$B) = StandardDist{$A}()
+
+        function Distributions.Distribution(m::PowerMeasure{$B})
+            StandardDist{$A}(map(dynamic ∘ length, m.axes)...)
+        end
+        Base.convert(::Type{Distribution}, m::PowerMeasure{$B}) = Distributions.Distribution(m)
+    end
+end
+
+@inline MeasureBase.to_origin(ν::StandardDist, y) = y
+@inline MeasureBase.from_origin(ν::StandardDist, x) = x
+
+
+@inline nonstddist(::StandardDist{D,0}) where {D} = D(Distributions.params(D())...)
+@inline function nonstddist(d::StandardDist{D,N}) where {D,N}
+    nonstd0 = nonstddist(StandardDist{D}())
+    reshape(Distributions.product_distribution(fill(nonstd0, length(d))), size(d))
+end
+
+
+(::Type{D})(d::StandardDist{D,0}) where {D<:Distribution{Univariate,Continuous}} = nonstddist(d)
+
+# TODO: Replace `fill` by `FillArrays.Fill` once Distributions fully supports this:
+(::Type{Distributions.Product})(d::StandardDist{D,1}) where {D} =
+    Distributions.Product(fill(StandardDist{D}(), length(d)))
+
+Base.convert(::Type{D}, d::StandardDist{D,0}) where {D<:Distribution{Univariate,Continuous}} = D(d)
+Base.convert(::Type{Distributions.Product}, d::StandardDist{D,1}) where {D} =
+    Distributions.Product(d)
+
+
+
+@inline Base.size(d::StandardDist) = d._size
+@inline Base.length(d::StandardDist) = prod(size(d))
+
+Base.eltype(::Type{<:StandardDist}) = Float64
+
+@inline Distributions.partype(d::StandardDist{D}) where {D} = Float64
+
+@inline StatsBase.params(d::StandardDist) = ()
+
+for f in (
+    :(Base.minimum),
+    :(Base.maximum),
+    :(Statistics.mean),
+    :(Statistics.median),
+    :(StatsBase.mode),
+    :(Statistics.var),
+    :(Statistics.std),
+    :(StatsBase.skewness),
+    :(StatsBase.kurtosis),
+    :(Distributions.location),
+    :(Distributions.scale),
+)
+    @eval begin
+        ($f)(d::StandardDist{D,0}) where {D} = ($f)(nonstddist(d))
+        ($f)(d::StandardDist{D,N}) where {D,N} = Fill(($f)(StandardDist{D}()), size(d)...)
+    end
+end
+
+StatsBase.modes(d::StandardDist) = [StatsBase.mode(d)]
+
+# ToDo: Define cov for N!=1?
+Statistics.cov(d::StandardDist{D,1}) where {D} = Diagonal(Statistics.var(d))
+Distributions.invcov(d::StandardDist{D,1}) where {D} =
+    Diagonal(Fill(inv(Statistics.var(StandardDist{D}())), length(d)))
+Distributions.logdetcov(d::StandardDist{D,1}) where {D} =
+    length(d) * log(Statistics.var(StandardDist{D}()))
+
+StatsBase.entropy(d::StandardDist{D,0}) where {D} = StatsBase.entropy(nonstddist(d))
+StatsBase.entropy(d::StandardDist{D,N}) where {D,N} =
+    length(d) * StatsBase.entropy(StandardDist{D}())
+
+
+Distributions.insupport(d::StandardDist{D,0}, x::Real) where {D} =
+    Distributions.insupport(nonstddist(d), x)
+
+function Distributions.insupport(d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N}
+    all(Base.Fix1(Distributions.insupport, StandardDist{D}()), checked_arg(d, x))
+end
+
+
+@inline Distributions.logpdf(d::StandardDist{D,0}, x::U) where {D,U} =
+    Distributions.logpdf(nonstddist(d), x)
+
+function Distributions.logpdf(d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N}
+    Distributions._logpdf(d, checked_arg(d, x))
+end
+
+# Explicit N=1/N=2 methods to avoid dispatch ambiguities with Distributions:
+function Distributions._logpdf(::StandardDist{D,1}, x::AbstractArray{<:Real,1}) where {D}
+    sum(Base.Fix1(Distributions.logpdf, StandardDist{D}()), x)
+end
+
+function Distributions._logpdf(::StandardDist{D,2}, x::AbstractArray{<:Real,2}) where {D}
+    sum(Base.Fix1(Distributions.logpdf, StandardDist{D}()), x)
+end
+
+function Distributions._logpdf(::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N}
+    sum(Base.Fix1(Distributions.logpdf, StandardDist{D}()), x)
+end
+
+
+Distributions.gradlogpdf(d::StandardDist{D,0}, x::Real) where {D} =
+    Distributions.gradlogpdf(nonstddist(d), x)
+
+function Distributions.gradlogpdf(d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N}
+    Distributions.gradlogpdf.(StandardDist{D}(), checked_arg(d, x))
+end
+
+
+# Explicit N=1/N=2 methods to avoid dispatch ambiguities with Distributions:
+function Distributions.pdf(d::StandardDist{D,1}, x::AbstractVector{U}) where {D,U<:Real}
+    Distributions._pdf(d, checked_arg(d, x))
+end
+
+function Distributions._pdf(d::StandardDist{D,1}, x::AbstractVector{U}) where {D,U<:Real}
+    exp(Distributions._logpdf(d, x))
+end
+
+function Distributions.pdf(d::StandardDist{D,2}, x::AbstractMatrix{U}) where {D,U<:Real}
+    Distributions._pdf(d, checked_arg(d, x))
+end
+
+function Distributions._pdf(d::StandardDist{D,2}, x::AbstractMatrix{U}) where {D,U<:Real}
+    exp(Distributions._logpdf(d, x))
+end
+
+function Distributions.pdf(d::StandardDist{D,N}, x::AbstractArray{U,N}) where {D,N,U<:Real}
+    Distributions._pdf(d, checked_arg(d, x))
+end
+
+function Distributions._pdf(d::StandardDist{D,N}, x::AbstractArray{U,N}) where {D,N,U<:Real}
+    exp(Distributions._logpdf(d, x))
+end
+
+
+for f in (
+    :(Distributions.logcdf),
+    :(Distributions.cdf),
+    :(Distributions.logccdf),
+    :(Distributions.ccdf),
+    :(Distributions.quantile),
+    :(Distributions.cquantile),
+    :(Distributions.invlogcdf),
+    :(Distributions.invlogccdf),
+    :(Distributions.mgf),
+    :(Distributions.cf),
+)
+    @eval begin
+        @inline ($f)(d::StandardDist, x::Real) = ($f)(nonstddist(d), x)
+    end
+end
+
+
+Base.rand(rng::AbstractRNG, d::StandardDist{D,0}) where {D} = rand(rng, nonstddist(d))
+Random.rand!(rng::AbstractRNG, d::StandardDist{D,0}, x::AbstractArray{<:Real,0}) where {D} =
+    (x[] = rand(rng, d); return x)
+Random.rand!(rng::AbstractRNG, d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N} =
+    rand!(rng, StandardDist{D}(), x)
+
+
+Distributions.truncated(d::StandardDist{D,0}, l::Real, u::Real) where {D} =
+    Distributions.truncated(nonstddist(d), l, u)
+
+Distributions.product_distribution(dists::AbstractVector{<:StandardDist{D,0}}) where {D} =
+    StandardDist{D}(size(dists)...)
+Distributions.product_distribution(dists::AbstractArray{<:StandardDist{D,0}}) where {D} =
+    StandardDist{D}(size(dists)...)
diff --git a/ext/MeasureBaseDistributionsExt/standard_normal.jl b/ext/MeasureBaseDistributionsExt/standard_normal.jl
new file mode 100644
index 00000000..6bc27d04
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/standard_normal.jl
@@ -0,0 +1,77 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+"""
+    const StandardNormal{N} = StandardDist{Normal,N}
+
+The standard normal distribution, scalar (`N == 0`) or as a product over an
+array of rank `N`.
+"""
+const StandardNormal{N} = StandardDist{Normal,N}
+export StandardNormal
+
+Distributions.Normal(d::StandardDist{Normal,0}) = Distributions.Normal()
+
+Distributions.MvNormal(d::StandardDist{Normal,1}) = MvNormal(PDMats.ScalMat(length(d), 1))
+Base.convert(::Type{Distributions.MvNormal}, d::StandardDist{Normal,1}) =
+    Distributions.MvNormal(d)
+
+Base.minimum(d::StandardDist{Normal,0}) = -Inf
+Base.maximum(d::StandardDist{Normal,0}) = +Inf
+
+Distributions.insupport(d::StandardDist{Normal,0}, x::Real) = !isnan(x)
+
+Distributions.location(d::StandardDist{Normal,0}) = Statistics.mean(d)
+Distributions.scale(d::StandardDist{Normal,0}) = Statistics.var(d)
+
+Statistics.mean(d::StandardDist{Normal,0}) = 0
+Statistics.mean(d::StandardDist{Normal,N}) where {N} = Zeros{Int}(size(d)...)
+
+Statistics.median(d::StandardDist{Normal}) = Statistics.mean(d)
+StatsBase.mode(d::StandardDist{Normal}) = Statistics.mean(d)
+
+StatsBase.modes(d::StandardDist{Normal,0}) = Zeros{Int}(1)
+
+Statistics.var(d::StandardDist{Normal,0}) = 1
+Statistics.var(d::StandardDist{Normal,N}) where {N} = Ones{Int}(size(d)...)
+
+Statistics.std(d::StandardDist{Normal,0}) = 1
+Statistics.std(d::StandardDist{Normal,N}) where {N} = Ones{Int}(size(d)...)
+
+StatsBase.skewness(d::StandardDist{Normal,0}) = 0
+StatsBase.kurtosis(d::StandardDist{Normal,0}) = 0
+
+StatsBase.entropy(d::StandardDist{Normal,0}) = muladd(log2π, 1 / 2, 1 / 2)
+
+Distributions.logpdf(d::StandardDist{Normal,0}, x::U) where {U<:Real} =
+    muladd(abs2(x), -U(1) / U(2), -log2π / U(2))
+Distributions.pdf(d::StandardDist{Normal,0}, x::U) where {U<:Real} =
+    invsqrt2π * exp(-abs2(x) / U(2))
+
+@inline Distributions.gradlogpdf(d::StandardDist{Normal,0}, x::Real) = -x
+
+@inline Distributions.logcdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normlogcdf(x)
+@inline Distributions.cdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normcdf(x)
+@inline Distributions.logccdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normlogccdf(x)
+@inline Distributions.ccdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normccdf(x)
+@inline Distributions.quantile(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvcdf(p)
+@inline Distributions.cquantile(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvccdf(p)
+@inline Distributions.invlogcdf(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvlogcdf(p)
+@inline Distributions.invlogccdf(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvlogccdf(p)
+
+Base.rand(rng::AbstractRNG, d::StandardDist{Normal,0}) = randn(rng)
+Base.rand(rng::AbstractRNG, d::StandardDist{Normal,N}) where {N} = randn(rng, size(d)...)
+Random.rand!(rng::AbstractRNG, d::StandardDist{Normal,N}, x::AbstractArray{<:Real,N}) where {N} =
+    Random.randn!(rng, x)
+
+Distributions.invcov(d::StandardDist{Normal,1}) = Distributions.cov(d)
+Distributions.logdetcov(d::StandardDist{Normal,1}) = 0
+
+
+function Distributions.sqmahal(d::StandardDist{Normal,N}, x::AbstractArray{<:Real,N}) where {N}
+    dot(x, checked_arg(d, x))
+end
+
+function Distributions.sqmahal!(r::AbstractVector, d::StandardDist{Normal,N}, x::AbstractMatrix) where {N}
+    x_cols = eachcol(checked_arg(d, first(eachcol(x))))
+    r .= dot.(x_cols, x_cols)
+end
diff --git a/ext/MeasureBaseDistributionsExt/standard_uniform.jl b/ext/MeasureBaseDistributionsExt/standard_uniform.jl
new file mode 100644
index 00000000..51398fcb
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/standard_uniform.jl
@@ -0,0 +1,79 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+"""
+    const StandardUniform{N} = StandardDist{Uniform,N}
+
+The standard uniform distribution, scalar (`N == 0`) or as a product over an
+array of rank `N`.
+"""
+const StandardUniform{N} = StandardDist{Uniform,N}
+export StandardUniform
+
+Distributions.Uniform(d::StandardDist{Uniform,0}) = Distributions.Uniform()
+
+Base.minimum(::StandardDist{Uniform,0}) = 0
+Base.maximum(::StandardDist{Uniform,0}) = 1
+
+Distributions.location(::StandardDist{Uniform,0}) = 0
+Distributions.scale(::StandardDist{Uniform,0}) = 1
+
+Statistics.mean(d::StandardDist{Uniform,0}) = 1 // 2
+Statistics.median(d::StandardDist{Uniform,0}) = Statistics.mean(d)
+StatsBase.mode(d::StandardDist{Uniform,0}) = Statistics.mean(d)
+StatsBase.modes(d::StandardDist{Uniform,0}) = Zeros{Int}(0)
+StatsBase.modes(d::StandardDist{Uniform,N}) where {N} = Fill(Zeros{Int}(size(d)))
+
+Statistics.var(d::StandardDist{Uniform,0}) = 1 // 12
+Statistics.std(d::StandardDist{Uniform,0}) = sqrt(Statistics.var(d))
+StatsBase.skewness(d::StandardDist{Uniform,0}) = 0
+StatsBase.kurtosis(d::StandardDist{Uniform,0}) = -6 // 5
+
+StatsBase.entropy(d::StandardDist{Uniform,0}) = 0
+
+
+function Distributions.logpdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real}
+    ifelse(Distributions.insupport(d, x), U(0), U(-Inf))
+end
+
+function Distributions.pdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real}
+    ifelse(Distributions.insupport(d, x), one(U), zero(U))
+end
+
+
+Distributions.logcdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} =
+    log(Distributions.cdf(d, x))
+
+function Distributions.cdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real}
+    ifelse(x < zero(U), zero(U), ifelse(x < one(U), x, one(U)))
+end
+
+Distributions.logccdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} =
+    log(Distributions.ccdf(d, x))
+
+Distributions.ccdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} =
+    one(x) - Distributions.cdf(d, x)
+
+
+function Distributions.quantile(d::StandardDist{Uniform,0}, p::U) where {U<:Real}
+    convert(float(U), p)
+end
+
+function Distributions.cquantile(d::StandardDist{Uniform,0}, p::U) where {U<:Real}
+    y = Distributions.quantile(d, p)
+    one(y) - y
+end
+
+
+Distributions.mgf(d::StandardDist{Uniform,0}, t::Real) = Distributions.mgf(nonstddist(d), t)
+Distributions.cf(d::StandardDist{Uniform,0}, t::Real) = Distributions.cf(nonstddist(d), t)
+
+Distributions.gradlogpdf(d::StandardDist{Uniform,0}, x::Real) = zero(x)
+
+function Distributions.gradlogpdf(d::StandardDist{Uniform,N}, x::AbstractArray{<:Real,N}) where {N}
+    zero(checked_arg(d, x))
+end
+
+Base.rand(rng::AbstractRNG, d::StandardDist{Uniform,0}) = rand(rng)
+Base.rand(rng::AbstractRNG, d::StandardDist{Uniform,N}) where {N} = rand(rng, size(d)...)
+Random.rand!(rng::AbstractRNG, d::StandardDist{Uniform,N}, x::AbstractArray{<:Real,N}) where {N} =
+    rand!(rng, x)
diff --git a/ext/MeasureBaseDistributionsExt/standardmv.jl b/ext/MeasureBaseDistributionsExt/standardmv.jl
new file mode 100644
index 00000000..c8e99039
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/standardmv.jl
@@ -0,0 +1,36 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+
+MeasureBase.getdof(d::AbstractMvNormal) = length(d)
+MeasureBase.getdof(m::AsMeasure{<:AbstractMvNormal}) = getdof(m.obj)
+
+MeasureBase.transport_origin(ν::MvNormal) = StandardDist{Normal}(length(ν))
+
+_cholesky_L(A) = cholesky(A).L
+_cholesky_L(A::Diagonal{<:Real}) = Diagonal(sqrt.(diag(A)))
+_cholesky_L(A::PDMats.PDiagMat{<:Real}) = Diagonal(sqrt.(A.diag))
+_cholesky_L(A::PDMats.ScalMat{<:Real}) = Diagonal(Fill(sqrt(A.value), A.dim))
+
+function MeasureBase.from_origin(ν::MvNormal, x)
+    A = _cholesky_L(ν.Σ)
+    b = ν.μ
+    muladd(A, x, b)
+end
+
+function MeasureBase.to_origin(ν::MvNormal, y)
+    A = _cholesky_L(ν.Σ)
+    b = ν.μ
+    A \ (y - b)
+end
+
+
+#DirichletMultinomial
+#Distributions.AbstractMvLogNormal
+#Distributions.AbstractMvTDist
+#Distributions.ProductDistribution{1}
+#Distributions.ReshapedDistribution{1, S, D} where {S<:ValueSupport, D<:(Distribution{<:ArrayLikeVariate, S})}
+#JointOrderStatistics
+#Multinomial
+#MultivariateMixture (alias for AbstractMixtureModel{ArrayLikeVariate{1}})
+#MvLogitNormal
+#VonMisesFisher
diff --git a/ext/MeasureBaseDistributionsExt/univariate.jl b/ext/MeasureBaseDistributionsExt/univariate.jl
new file mode 100644
index 00000000..126fe36d
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/univariate.jl
@@ -0,0 +1,142 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+
+@inline MeasureBase.getdof(::Distribution{Univariate}) = static(1)
+
+@inline MeasureBase.check_dof(a::Distribution{Univariate}, b::Distribution{Univariate}) = nothing
+
+
+# Generic transformations to/from StdUniform via cdf/quantile:
+
+
+_dist_params_numtype(d::Distribution) = real_numtype(typeof(Distributions.params(d)))
+
+
+@inline _trafo_cdf(d::Distribution{Univariate,Continuous}, x::Real) =
+    _trafo_cdf_impl(_dist_params_numtype(d), d, x)
+
+@inline _trafo_cdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Real) =
+    Distributions.cdf(d, x)
+
+
+@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, u::Real) =
+    _trafo_quantile_impl(_dist_params_numtype(d), d, u)
+
+@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, u::Real) =
+    _trafo_quantile_impl_generic(d, u)
+
+
+@inline _trafo_quantile_impl_generic(d::Distribution{Univariate,Continuous}, u::Real) =
+    Distributions.quantile(d, u)
+
+# Workaround for Beta dist, current quantile implementation only supports Float64:
+@inline function _trafo_quantile_impl_generic(d::Beta{T}, u::Union{Integer,AbstractFloat}) where {T<:Union{Integer,AbstractFloat}}
+    Distributions.quantile(d, convert(promote_type(Float64, typeof(u)), u))
+end
+
+# Workaround for rounding errors that can result in quantile values outside of support of Truncated:
+@inline function _trafo_quantile_impl_generic(d::Distributions.Truncated{<:Distribution{Univariate,Continuous}}, u::Real)
+    x = Distributions.quantile(d, u)
+    T = typeof(x)
+    min_x = T(minimum(d))
+    max_x = T(maximum(d))
+    if x < min_x && isapprox(x, min_x, atol = 4 * eps(T))
+        min_x
+    elseif x > max_x && isapprox(x, max_x, atol = 4 * eps(T))
+        max_x
+    else
+        x
+    end
+end
+
+
+@inline function _result_numtype(d::Distribution{Univariate}, x::T) where {T<:Real}
+    float(promote_type(T, _dist_params_numtype(d)))
+end
+
+
+@inline function MeasureBase.transport_def(::StdUniform, μ::Distribution{Univariate,Continuous}, x)
+    R = _result_numtype(μ, x)
+    if Distributions.insupport(μ, x)
+        y = _trafo_cdf(μ, x)
+        convert(R, y)
+    else
+        convert(R, NaN)
+    end
+end
+
+
+@inline function MeasureBase.transport_def(ν::Distribution{Univariate,Continuous}, ::StdUniform, x::T) where {T}
+    R = _result_numtype(ν, x)
+    TF = float(T)
+    if 0 <= x <= 1
+        # Avoid x ≈ 0 and x ≈ 1 to avoid infinite variate values for target distributions with infinite support:
+        mod_x = ifelse(x ≈ 0, zero(TF) + eps(TF), ifelse(x ≈ 1, one(TF) - eps(TF), convert(TF, x)))
+        y = _trafo_quantile(ν, mod_x)
+        convert(R, y)
+    else
+        convert(R, NaN)
+    end
+end
+
+
+# Use standard measures as transformation origin for scaled/translated equivalents:
+
+function _origin_to_affine(ν::Distribution{Univariate}, y::T) where {T<:Real}
+    trg_offs, trg_scale = Distributions.location(ν), Distributions.scale(ν)
+    x = muladd(y, trg_scale, trg_offs)
+    convert(_result_numtype(ν, y), x)
+end
+
+function _affine_to_origin(μ::Distribution{Univariate}, x::T) where {T<:Real}
+    src_offs, src_scale = Distributions.location(μ), Distributions.scale(μ)
+    y = (x - src_offs) / src_scale
+    convert(_result_numtype(μ, x), y)
+end
+
+for (A, B) in [
+    (Uniform, StdUniform),
+    (Logistic, StdLogistic),
+    (Normal, StdNormal)
+]
+    @eval begin
+        @inline MeasureBase.transport_origin(::$A) = $B()
+        @inline MeasureBase.to_origin(ν::$A, y) = _affine_to_origin(ν, y)
+        @inline MeasureBase.from_origin(ν::$A, x) = _origin_to_affine(ν, x)
+    end
+end
+
+@inline MeasureBase.transport_origin(::Exponential) = StdExponential()
+@inline MeasureBase.to_origin(ν::Exponential, y) = Distributions.scale(ν) \ y
+@inline MeasureBase.from_origin(ν::Exponential, x) = Distributions.scale(ν) * x
+
+
+# Use the underlying distribution as transformation origin for affine
+# transformed distributions:
+
+@inline MeasureBase.transport_origin(d::Distributions.AffineDistribution) = d.ρ
+@inline MeasureBase.from_origin(d::Distributions.AffineDistribution, x) = muladd(d.σ, x, d.μ)
+@inline MeasureBase.to_origin(d::Distributions.AffineDistribution, y) = d.σ \ (y - d.μ)
+
+
+
+# Transform between univariate and single-element power measure
+
+function MeasureBase.transport_def(ν::Distribution{Univariate}, μ::PowerMeasure{<:StdMeasure}, x)
+    return transport_def(ν, μ.parent, only(x))
+end
+
+function MeasureBase.transport_def(ν::PowerMeasure{<:StdMeasure}, μ::Distribution{Univariate}, x)
+    return Fill(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)...)
+end
+
+
+# Transform between univariate and single-element standard multivariate
+
+function MeasureBase.transport_def(ν::Distribution{Univariate}, μ::StandardDist{D,1}, x) where {D}
+    return transport_def(ν, StandardDist{D}(), only(x))
+end
+
+function MeasureBase.transport_def(ν::StandardDist{D,1}, μ::Distribution{Univariate}, x) where {D}
+    return Fill(transport_def(StandardDist{D}(), μ, only(x)), size(ν)...)
+end
diff --git a/ext/MeasureBaseDistributionsForwardDiffExt.jl b/ext/MeasureBaseDistributionsForwardDiffExt.jl
index 36218eec..245c8e60 100644
--- a/ext/MeasureBaseDistributionsForwardDiffExt.jl
+++ b/ext/MeasureBaseDistributionsForwardDiffExt.jl
@@ -2,8 +2,35 @@
 
 module MeasureBaseDistributionsForwardDiffExt
 
-using MeasureBase
+import MeasureBase
 import Distributions
 import ForwardDiff
 
+using Distributions: Distribution, Univariate, Continuous, Beta
+
+@inline function MeasureBase._trafo_cdf_impl(
+    ::Type{<:Union{Integer,AbstractFloat}},
+    d::Distribution{Univariate,Continuous},
+    x::ForwardDiff.Dual{TAG},
+) where {TAG}
+    x_v = ForwardDiff.value(x)
+    u = Distributions.cdf(d, x_v)
+    dudx = Distributions.pdf(d, x_v)
+    ForwardDiff.Dual{TAG}(u, dudx * ForwardDiff.partials(x))
+end
+
+@inline function MeasureBase._trafo_quantile_impl(
+    ::Type{<:Union{Integer,AbstractFloat}},
+    d::Distribution{Univariate,Continuous},
+    u::ForwardDiff.Dual{TAG},
+) where {TAG}
+    x = MeasureBase._trafo_quantile_impl_generic(d, ForwardDiff.value(u))
+    dxdu = inv(Distributions.pdf(d, x))
+    ForwardDiff.Dual{TAG}(x, dxdu * ForwardDiff.partials(u))
+end
+
+# Workaround for Beta dist, ForwardDiff doesn't work for parameters:
+@inline MeasureBase._trafo_quantile_impl_generic(d::Beta{T}, u::Real) where {T<:ForwardDiff.Dual} =
+    convert(float(typeof(u)), NaN)
+
 end # module MeasureBaseDistributionsForwardDiffExt
diff --git a/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl b/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl
new file mode 100644
index 00000000..667af023
--- /dev/null
+++ b/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl
@@ -0,0 +1,25 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseDistributionsForwardDiffPullbacksExt
+
+import MeasureBase
+using MeasureBase: StdMeasure, transport_def
+
+import Distributions
+using Distributions: Distribution, Univariate
+
+import ChainRulesCore
+using ForwardDiffPullbacks: fwddiff
+
+# Use ForwardDiff for univariate transformations:
+@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::Distribution{Univariate}, μ::Distribution{Univariate}, x::Any)
+    ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x)
+end
+@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::StdMeasure, μ::Distribution{Univariate}, x::Any)
+    ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x)
+end
+@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::Distribution{Univariate}, μ::StdMeasure, x::Any)
+    ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x)
+end
+
+end # module MeasureBaseDistributionsForwardDiffPullbacksExt
diff --git a/ext/MeasureBaseDistributionsMooncakeExt.jl b/ext/MeasureBaseDistributionsMooncakeExt.jl
new file mode 100644
index 00000000..75bde91b
--- /dev/null
+++ b/ext/MeasureBaseDistributionsMooncakeExt.jl
@@ -0,0 +1,21 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseDistributionsMooncakeExt
+
+using MeasureBase
+import Distributions
+import Mooncake
+using Mooncake: @zero_derivative, MinimalCtx
+
+using Distributions: Distribution
+using MeasureBase: _dist_params_numtype
+
+# The distribution transports themselves need no rules here: Mooncake
+# provides rules for Distributions and StatsFuns/SpecialFunctions, so it
+# differentiates the cdf/quantile-based transports natively. The
+# ForwardDiffPullbacks-based rules for `transport_def` are a
+# Zygote/ChainRules pathway.
+
+@zero_derivative MinimalCtx Tuple{typeof(_dist_params_numtype),Distribution}
+
+end # module MeasureBaseDistributionsMooncakeExt
diff --git a/ext/MeasureBaseForwardDiffExt.jl b/ext/MeasureBaseForwardDiffExt.jl
index 8a1cab44..20113c7f 100644
--- a/ext/MeasureBaseForwardDiffExt.jl
+++ b/ext/MeasureBaseForwardDiffExt.jl
@@ -3,6 +3,7 @@
 module MeasureBaseForwardDiffExt
 
 using MeasureBase
+using MeasureBase: containsnan, firsttype
 import ForwardDiff
 
 function MeasureBase.containsnan(x::ForwardDiff.Dual)
@@ -11,4 +12,7 @@ function MeasureBase.containsnan(x::ForwardDiff.Dual)
     return a || b
 end
 
+MeasureBase.firsttype(::Type{T}, ::Type{<:ForwardDiff.Dual{tag,<:Real,N}}) where {T<:Real,tag,N} =
+    ForwardDiff.Dual{tag,T,N}
+
 end # module MeasureBaseForwardDiffExt
diff --git a/ext/MeasureBaseForwardDiffPullbacksExt.jl b/ext/MeasureBaseForwardDiffPullbacksExt.jl
new file mode 100644
index 00000000..72ffc751
--- /dev/null
+++ b/ext/MeasureBaseForwardDiffPullbacksExt.jl
@@ -0,0 +1,10 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseForwardDiffPullbacksExt
+
+import MeasureBase
+using ForwardDiffPullbacks: fwddiff
+
+MeasureBase._fwddiff(f::Function) = fwddiff(f)
+
+end # module MeasureBaseForwardDiffPullbacksExt
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index 099ee806..aa7b6e20 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -65,37 +65,24 @@ function Base.:+(μ::AbstractMeasure, ν::AbstractMeasure)
     superpose(μ, ν)
 end
 
-oneplus(x::ULogarithmic) = exp(ULogarithmic, log1pexp(x.log))
-
-@inline function density_def(s::SuperpositionMeasure{Tuple{A,B}}, x) where {A,B}
-    (μ, ν) = s.components
-
-    istrue(insupport(μ, x)) || return exp(ULogarithmic, logdensity_def(ν, x))
-    istrue(insupport(ν, x)) || return exp(ULogarithmic, logdensity_def(μ, x))
-
-    α = basemeasure(μ)
-    β = basemeasure(ν)
-    dμ_dα = exp(ULogarithmic, logdensity_def(μ, x))
-    dν_dβ = exp(ULogarithmic, logdensity_def(ν, x))
-    dα_dβ = exp(ULogarithmic, logdensity_rel(α, β, x))
-    dβ_dα = inv(dα_dβ)
-    return dμ_dα / oneplus(dβ_dα) + dν_dβ / oneplus(dα_dβ)
-end
+@inline _ulogexp(x) = exp(ULogarithmic, dynamic(x))
 
 function density_def(s::SuperpositionMeasure, x)
-    T = typeof(s)
-    msg = """
-    Not implemented: There is no method
-    density_def(::$T, x)
-    """
-    error(msg)
+    cs = values(s.components)
+    αs = map(basemeasure, cs)
+    idxs = eachindex(cs)
+    sum(idxs) do i
+        dμᵢ_dαᵢ = _ulogexp(logdensity_def(cs[i], x))
+        istrue(insupport(cs[i], x)) || return zero(dμᵢ_dαᵢ)
+        dΣα_dαᵢ = sum(idxs) do j
+            dαⱼ_dαᵢ = _ulogexp(logdensity_rel(αs[j], αs[i], x))
+            istrue(insupport(cs[j], x)) ? dαⱼ_dαᵢ : zero(dαⱼ_dαᵢ)
+        end
+        dμᵢ_dαᵢ / dΣα_dαᵢ
+    end
 end
 
-@inline function logdensity_def(
-    μ::T,
-    ν::T,
-    x,
-) where {T<:(SuperpositionMeasure{Tuple{A,B}} where {A,B})}
+@inline function logdensity_def(μ::T, ν::T, x) where {T<:SuperpositionMeasure}
     if μ === ν
         return zero(return_type(logdensity_def, (μ, x)))
     else
@@ -103,44 +90,49 @@ end
     end
 end
 
-@inline function logdensity_def(
-    s::T,
-    β,
-    x,
-) where {T<:(SuperpositionMeasure{Tuple{A,B}} where {A,B})}
-    (μ, ν) = s.components
-
-    istrue(insupport(μ, x)) || return logdensity_rel(ν, β, x)
-    istrue(insupport(ν, x)) || return logdensity_rel(μ, β, x)
-    return logaddexp(logdensity_rel(μ, β, x), logdensity_rel(ν, β, x))
+function _superpos_logdensity_rel(s::SuperpositionMeasure, β, x)
+    cs = values(s.components)
+    ds = map(cs) do μ
+        istrue(insupport(μ, x)) ? dynamic(logdensity_rel(μ, β, x)) : -Inf
+    end
+    logsumexp(ds)
 end
 
-@inline function logdensity_def(
-    s::SuperpositionMeasure{Tuple{A,B}},
-    β::SuperpositionMeasure,
-    x,
-) where {A,B}
-    (μ, ν) = s.components
-    istrue(insupport(μ, x)) || return logdensity_rel(ν, β, x)
-    istrue(insupport(ν, x)) || return logdensity_rel(μ, β, x)
-    return logaddexp(logdensity_rel(μ, β, x), logdensity_rel(ν, β, x))
-end
+@inline logdensity_def(s::SuperpositionMeasure, β, x) = _superpos_logdensity_rel(s, β, x)
 
-@inline function logdensity_def(s, β::(SuperpositionMeasure{Tuple{A,B}} where {A,B}), x)
-    -logdensity_def(β, s, x)
-end
+@inline logdensity_def(s::SuperpositionMeasure, β::SuperpositionMeasure, x) =
+    _superpos_logdensity_rel(s, β, x)
+
+@inline logdensity_def(s, β::SuperpositionMeasure, x) = -_superpos_logdensity_rel(β, s, x)
 
 @inline logdensity_def(s::SuperpositionMeasure, x) = log(density_def(s, x))
 
-function basemeasure(μ::SuperpositionMeasure{Tuple{A,B}}) where {A,B}
+function basemeasure(μ::SuperpositionMeasure{<:Tuple})
     superpose(map(basemeasure, μ.components)...)
 end
+
+function basemeasure(μ::SuperpositionMeasure{<:AbstractArray})
+    bases = map(basemeasure, μ.components)
+    allequal(bases) ? weightedmeasure(log(length(bases)), first(bases)) : superpose(bases)
+end
+
 basemeasure(μ::SuperpositionMeasure) = superpose(map(basemeasure, μ.components))
 
-# TODO: Fix `rand` method (this one is wrong)
-# function Base.rand(μ::SuperpositionMeasure{X,N}) where {X,N}
-#     return rand(rand(μ.components))
-# end
+function Base.rand(rng::AbstractRNG, ::Type{T}, μ::SuperpositionMeasure) where {T}
+    components = values(μ.components)
+    masses = map(massof, components)
+    total = sum(masses)
+    total isa AbstractUnknownMass && throw(
+        ArgumentError("Cannot sample from a superposition of measures of unknown mass"),
+    )
+    threshold = rand(rng) * dynamic(total)
+    csum = zero(threshold)
+    for (mass, c) in zip(masses, components)
+        csum += dynamic(mass)
+        csum >= threshold && return rand(rng, T, c)
+    end
+    return rand(rng, T, last(components))
+end
 
 @inline function insupport(d::SuperpositionMeasure, x)
     any(d.components) do c
diff --git a/src/utils.jl b/src/utils.jl
index 5d05d8b1..c1e97034 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -185,3 +185,54 @@ isapproxzero(A::AbstractArray) = all(isapproxzero, A)
 
 isapproxone(x::T) where {T<:Real} = x ≈ one(T)
 isapproxone(A::AbstractArray) = all(isapproxone, A)
+
+containsnan(x::Real) = isnan(x)
+containsnan(x) = any(containsnan, x)
+
+
+# ForwardDiffPullbacks dummy `fwddiff`, overloaded by
+# ForwardDiffPullbacks extension when loaded:
+@inline _fwddiff(f) = f
+
+
+# Autodiff ignore:
+
+@inline _adignore_call(f) = f()
+
+macro _adignore(expr)
+    :(_adignore_call(() -> $(esc(expr))))
+end
+
+
+"""
+    MeasureBase.convert_realtype(::Type{T}, x) where {T<:Real}
+
+Convert `x` to use `T` as its underlying type for real numbers.
+"""
+function convert_realtype end
+
+@inline convert_realtype(::Type{T}, x::T) where {T<:Real} = x
+@inline convert_realtype(::Type{T}, x::AbstractArray{T}) where {T<:Real} = x
+@inline convert_realtype(::Type{T}, x::U) where {T<:Real,U<:Real} = T(x)
+convert_realtype(::Type{T}, x::AbstractArray{U}) where {T<:Real,U<:Real} = T.(x)
+convert_realtype(::Type{T}, x::Union{Tuple,NamedTuple}) where {T<:Real} =
+    map(Base.Fix1(convert_realtype, T), x)
+convert_realtype(::Type{T}, x::AbstractArray) where {T<:Real} =
+    map(Base.Fix1(convert_realtype, T), x)
+
+"""
+    MeasureBase.firsttype(::Type{T}, ::Type{U}) where {T<:Real,U<:Real}
+
+Return the first type, but as a dual number type if the second one is dual.
+"""
+function firsttype end
+
+firsttype(::Type{T}, ::Type{U}) where {T<:Real,U<:Real} = T
+
+
+
+# Distributions implementation hooks:
+function _trafo_cdf_impl end
+function _trafo_quantile_impl end
+function _trafo_quantile_impl_generic end
+function _dist_params_numtype end
diff --git a/test/Project.toml b/test/Project.toml
index fec229fd..cb68fc93 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,22 +1,29 @@
 [deps]
 AffineMaps = "2c83c9a8-abf5-4329-a0d7-deffaf474661"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 PropertyFunctions = "09e99361-2bb8-48a2-a80f-de58f0739eb4"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
diff --git a/test/distributions/getjacobian.jl b/test/distributions/getjacobian.jl
new file mode 100644
index 00000000..87de7b86
--- /dev/null
+++ b/test/distributions/getjacobian.jl
@@ -0,0 +1,34 @@
+# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT).
+
+import ForwardDiff
+
+torv_and_back(V::AbstractVector{<:Real}) = V, identity
+torv_and_back(x::Real) = [x], V -> V[1]
+torv_and_back(x::Complex) = [real(x), imag(x)], V -> Complex(V[1], V[2])
+torv_and_back(x::NTuple{N}) where N = [x...], V -> ntuple(i -> V[i], Val(N))
+
+function torv_and_back(x::Ref)
+    xval = x[]
+    V, to_xval = torv_and_back(xval)
+    back_to_ref(V) = Ref(to_xval(V))
+    return (V, back_to_ref)
+end
+
+torv_and_back(A::AbstractArray{<:Real}) = vec(A), V -> reshape(V, size(A))
+
+function torv_and_back(A::AbstractArray{Complex{T}, N}) where {T<:Real, N}
+    RA = cat(real.(A), imag.(A), dims = N+1) 
+    V, to_array = torv_and_back(RA)
+    function back_to_complex(V)
+        RA = to_array(V)
+        Complex.(view(RA, map(_ -> :, size(A))..., 1), view(RA, map(_ -> :, size(A))..., 2))
+    end
+    return (V, back_to_complex)
+end
+
+
+function getjacobian(f, x)
+    V, to_x = torv_and_back(x)
+    vf(V) = torv_and_back(f(to_x(V)))[1]
+    ForwardDiff.jacobian(vf, V)
+end
diff --git a/test/distributions/test_autodiff_utils.jl b/test/distributions/test_autodiff_utils.jl
new file mode 100644
index 00000000..5197725b
--- /dev/null
+++ b/test/distributions/test_autodiff_utils.jl
@@ -0,0 +1,18 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using LinearAlgebra
+using Distributions, ArraysOfArrays
+import ForwardDiff, Zygote
+
+
+@testset "trafo_utils" begin
+    xs = rand(5)
+    @test Zygote.jacobian(MeasureBase._pushfront, xs, 42)[1] ≈ ForwardDiff.jacobian(xs -> MeasureBase._pushfront(xs, 1), xs)
+    @test Zygote.jacobian(MeasureBase._pushfront, xs, 42)[2] ≈ vec(ForwardDiff.jacobian(x -> MeasureBase._pushfront(xs, x[1]), [42]))
+    @test Zygote.jacobian(MeasureBase._pushback, xs, 42)[1] ≈ ForwardDiff.jacobian(xs -> MeasureBase._pushback(xs, 1), xs)
+    @test Zygote.jacobian(MeasureBase._pushback, xs, 42)[2] ≈ vec(ForwardDiff.jacobian(x -> MeasureBase._pushback(xs, x[1]), [42]))
+    @test Zygote.jacobian(MeasureBase._rev_cumsum, xs)[1] ≈ ForwardDiff.jacobian(MeasureBase._rev_cumsum, xs)
+    @test Zygote.jacobian(MeasureBase._exp_cumsum_log, xs)[1] ≈ ForwardDiff.jacobian(MeasureBase._exp_cumsum_log, xs) ≈ ForwardDiff.jacobian(cumprod, xs)
+end
diff --git a/test/distributions/test_conversions.jl b/test/distributions/test_conversions.jl
new file mode 100644
index 00000000..0f9c4d80
--- /dev/null
+++ b/test/distributions/test_conversions.jl
@@ -0,0 +1,113 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using Random, Statistics, LinearAlgebra
+using Distributions
+using StableRNGs
+
+import MeasureBase
+using MeasureBase: AbstractMeasure, AsMeasure, asmeasure
+using MeasureBase: StdUniform, StdNormal, StdExponential, StdLogistic
+using MeasureBase: SuperpositionMeasure, PushforwardMeasure, ProductMeasure
+using MeasureBase: logdensityof, massof, insupport
+
+
+@testset "conversions" begin
+    stblrng() = StableRNG(789990641)
+
+    function test_conversion(d::Distribution, ::Type{M}) where {M}
+        @testset "conversion $(typeof(d).name) <-> $M" begin
+            m = asmeasure(d)
+            @test m isa M
+            @test typeof(convert(AbstractMeasure, d)) === typeof(m)
+            @test_throws ArgumentError AsMeasure{typeof(d)}(d)
+
+            d2 = convert(Distribution, m)
+            @test d2 isa Distribution
+            @test typeof(Distributions.Distribution(m)) === typeof(d2)
+
+            for x in (rand(stblrng(), d) for _ in 1:10)
+                @test logdensityof(m, x) ≈ logpdf(d, x)
+                @test logpdf(d2, x) ≈ logpdf(d, x)
+                @test insupport(m, x)
+            end
+
+            x = rand(stblrng(), Float64, m)
+            # Tuple-marginal product measures have tuple variates:
+            x isa Tuple ? (@test length(x) == length(d)) : (@test size(x) == size(d))
+            @test insupport(m, x)
+        end
+    end
+
+    @testset "Dirac" begin
+        d = Distributions.Dirac(4.2)
+        m = @inferred asmeasure(d)
+        @test m === MeasureBase.Dirac(4.2)
+        @test_throws ArgumentError AsMeasure{typeof(d)}(d)
+        @test @inferred(Distributions.Distribution(m)) === d
+    end
+
+    @testset "products" begin
+        test_conversion(product_distribution(Weibull.([0.7, 1.1, 1.3])), ProductMeasure)
+        test_conversion(product_distribution(Poisson.([0.7, 1.4])), ProductMeasure)
+
+        if isdefined(Distributions, :ProductDistribution)
+            test_conversion(product_distribution(Weibull(0.7), Exponential(1.3)), ProductMeasure)
+        end
+    end
+
+    @testset "reshaped" begin
+        test_conversion(reshape(MvNormal([0.7, 0.9], [1.4 0.5; 0.5 1.1]), 1, 2), PushforwardMeasure)
+        test_conversion(reshape(product_distribution(Weibull.([0.7, 1.1, 1.3, 0.9, 1.2, 0.8])), 2, 3), PushforwardMeasure)
+    end
+
+    @testset "mixtures" begin
+        test_conversion(MixtureModel([Normal(-1.0, 1.0), Normal(2.0, 3.0)], [0.3, 0.7]), SuperpositionMeasure)
+        test_conversion(MixtureModel([Normal(-2.0, 1.0), Normal(0.0, 2.0), Normal(3.0, 1.0)], [0.2, 0.5, 0.3]), SuperpositionMeasure)
+        test_conversion(MixtureModel([Exponential(0.3), Weibull(2.0, 1.0)], [0.4, 0.6]), SuperpositionMeasure)
+        test_conversion(MixtureModel([MvNormal([0.0, 0.0], I(2)), MvNormal([2.0, 2.0], 2 * I(2))], [0.3, 0.7]), SuperpositionMeasure)
+        test_conversion(UnivariateGMM([-1.0, 2.0], [1.0, 0.5], Categorical([0.4, 0.6])), SuperpositionMeasure)
+
+        d = MixtureModel([Normal(-1.0, 1.0), Normal(2.0, 3.0)], [0.3, 0.7])
+        m = asmeasure(d)
+        @test probs(convert(Distribution, m)) ≈ probs(d)
+        @test massof(m) ≈ 1
+        @test massof(asmeasure(Normal())) == 1
+        @test mean(rand(stblrng(), Float64, m^1000)) ≈ mean(d) atol = 0.3
+
+        # Hand-built superpositions of weighted probability measures behave
+        # like mixtures:
+        m2 = 0.3 * asmeasure(Normal(-1.0, 1.0)) + 0.7 * asmeasure(Normal(2.0, 3.0))
+        for x in (rand(stblrng(), d) for _ in 1:10)
+            @test logdensityof(m2, x) ≈ logpdf(d, x)
+        end
+    end
+
+    @testset "standard distributions" begin
+        @test StandardUniform === StandardDist{Uniform}
+        @test StandardNormal === StandardDist{Normal}
+        @test StandardUniform{0} === StandardDist{Uniform,0}
+        @test StandardNormal{1} === StandardDist{Normal,1}
+
+        for (D, B) in [
+            (Uniform, StdUniform()),
+            (Exponential, StdExponential()),
+            (Logistic, StdLogistic()),
+            (Normal, StdNormal()),
+        ]
+            @test @inferred(asmeasure(StandardDist{D}())) === B
+            @test @inferred(asmeasure(StandardDist{D}(3))) == B^3
+            @test @inferred(asmeasure(StandardDist{D}(2, 3))) == B^(2, 3)
+
+            @test @inferred(Distributions.Distribution(B)) === StandardDist{D}()
+            @test @inferred(convert(Distribution, B)) === StandardDist{D}()
+            @test @inferred(Distributions.Distribution(B^3)) == StandardDist{D}(3)
+            @test @inferred(convert(Distribution, B^(2, 3))) == StandardDist{D}(2, 3)
+
+            d = StandardDist{D}(3)
+            x = rand(stblrng(), d)
+            @test logdensityof(asmeasure(d), x) ≈ logpdf(d, x)
+        end
+    end
+end
diff --git a/test/distributions/test_distribution_measure.jl b/test/distributions/test_distribution_measure.jl
new file mode 100644
index 00000000..7715e58a
--- /dev/null
+++ b/test/distributions/test_distribution_measure.jl
@@ -0,0 +1,53 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+import Distributions
+using Distributions: Distribution
+import MeasureBase
+using MeasureBase: AbstractMeasure
+
+@testset "Measure interface" begin
+    d = Distributions.Weibull()
+    @test @inferred(AbstractMeasure(d)) isa AbstractMeasure
+    @test @inferred(AbstractMeasure(d)) isa DistributionMeasure
+    @test @inferred(convert(AbstractMeasure, d)) isa AbstractMeasure
+    @test @inferred(convert(AbstractMeasure, d)) isa DistributionMeasure
+    @test @inferred(Distribution(AbstractMeasure(d))) === d
+    @test @inferred(convert(Distribution, convert(AbstractMeasure, d))) === d
+
+
+    c0 = AbstractMeasure(Distributions.Weibull(0.7, 1.3))
+    c1 = AbstractMeasure(Distributions.MvNormal([0.7, 0.9], [1.4 0.5; 0.5 1.1]))
+
+    d0 = AbstractMeasure(Distributions.Poisson(0.7))
+    d1 = AbstractMeasure(Distributions.product_distribution(Distributions.Poisson.([0.7, 1.4])))
+
+    for μ in [c0, c1, d0, d1]
+        d = Distribution(μ)
+        x = rand(μ)
+        @test @inferred(MeasureBase.logdensity_def(μ, x)) == Distributions.logpdf(d, x)
+        @test @inferred(MeasureBase.unsafe_logdensityof(μ, x)) == Distributions.logpdf(d, x)
+
+        MeasureBase.Interface.test_interface(d)
+    end
+
+    @test @inferred(MeasureBase.basemeasure(c0)) == MeasureBase.Lebesgue(MeasureBase.ℝ)
+    @test @inferred(MeasureBase.basemeasure(c1)) == MeasureBase.Lebesgue(MeasureBase.ℝ) ^ 2
+
+    @test @inferred(MeasureBase.insupport(c0, 3)) == true
+    @test @inferred(MeasureBase.insupport(c0, -3)) == false
+    @test @inferred(MeasureBase.insupport(c1, [0.1, 0.2])) == true
+    @test @inferred(MeasureBase.insupport(d0, 3)) == true
+    @test @inferred(MeasureBase.insupport(d0, 3.2)) == false
+    @test @inferred(MeasureBase.insupport(d1, [1, 2])) == true
+    @test @inferred(MeasureBase.insupport(d1, [1.1, 2.2])) == false
+
+    @test MeasureBase.paramnames(c0) == (:α, :θ)
+    if VERSION >= v"1.8"
+        @test @inferred(MeasureBase.params(c0)) == (α = 0.7, θ = 1.3)
+    else
+        # v1.6 can't type-infer this:
+        @test (MeasureBase.params(c0)) == (α = 0.7, θ = 1.3)
+    end
+end
diff --git a/test/distributions/test_distributions.jl b/test/distributions/test_distributions.jl
new file mode 100644
index 00000000..65fab330
--- /dev/null
+++ b/test/distributions/test_distributions.jl
@@ -0,0 +1,24 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+using MeasureBase
+using Distributions
+import ForwardDiff, ForwardDiffPullbacks, ChainRulesCore
+
+const MeasureBaseDistributionsExt = Base.get_extension(MeasureBase, :MeasureBaseDistributionsExt)
+@test MeasureBaseDistributionsExt isa Module
+
+using .MeasureBaseDistributionsExt:
+    StandardDist, StandardUniform, StandardNormal, DistributionMeasure, nonstddist
+
+@testset "Distributions extension" begin
+    include("test_autodiff_utils.jl")
+    include("test_measure_interface.jl")
+    include("test_distribution_measure.jl")
+    include("test_standard_dist.jl")
+    include("test_standard_uniform.jl")
+    include("test_standard_normal.jl")
+    include("test_conversions.jl")
+    include("test_transport.jl")
+    include("test_mooncake.jl")
+end
diff --git a/test/distributions/test_measure_interface.jl b/test/distributions/test_measure_interface.jl
new file mode 100644
index 00000000..d11d2889
--- /dev/null
+++ b/test/distributions/test_measure_interface.jl
@@ -0,0 +1,43 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+import Distributions
+import MeasureBase
+
+@testset "Measure interface" begin
+    c0 = Distributions.Weibull(0.7, 1.3)
+    c1 = Distributions.MvNormal([0.7, 0.9], [1.4 0.5; 0.5 1.1])
+
+    d0 = Distributions.Poisson(0.7)
+    d1 = Distributions.product_distribution(Distributions.Poisson.([0.7, 1.4]))
+
+    for d in [c0, c1, d0, d1]
+        x = rand(d)
+        @test @inferred(MeasureBase.logdensity_def(d, x)) == Distributions.logpdf(d, x)
+        @test @inferred(MeasureBase.unsafe_logdensityof(d, x)) == Distributions.logpdf(d, x)
+
+        MeasureBase.Interface.test_interface(d)
+    end
+
+    @test @inferred(MeasureBase.basemeasure(c0)) == MeasureBase.Lebesgue(MeasureBase.ℝ)
+    @test @inferred(MeasureBase.basemeasure(c1)) == MeasureBase.Lebesgue(MeasureBase.ℝ) ^ 2
+
+    @test @inferred(MeasureBase.insupport(c0, 3)) == true
+    @test @inferred(MeasureBase.insupport(c0, -3)) == false
+    @test @inferred(MeasureBase.insupport(c1, [0.1, 0.2])) == true
+    @test @inferred(MeasureBase.insupport(d0, 3)) == true
+    @test @inferred(MeasureBase.insupport(d0, 3.2)) == false
+    @test @inferred(MeasureBase.insupport(d1, [1, 2])) == true
+    @test @inferred(MeasureBase.insupport(d1, [1.1, 2.2])) == false
+
+    @test MeasureBase.paramnames(c0) == (:α, :θ)
+    if VERSION >= v"1.8"
+        @test @inferred(MeasureBase.params(c0)) == (α = 0.7, θ = 1.3)
+    else
+        # v1.6 can't type-infer this:
+        @test (MeasureBase.params(c0)) == (α = 0.7, θ = 1.3)
+    end
+
+    @test MeasureBase.∫(x -> Distributions.Normal(x, 0), Distributions.Normal()) isa MeasureBase.DensityMeasure
+end
diff --git a/test/distributions/test_mooncake.jl b/test/distributions/test_mooncake.jl
new file mode 100644
index 00000000..4a20e46c
--- /dev/null
+++ b/test/distributions/test_mooncake.jl
@@ -0,0 +1,68 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using Random, LinearAlgebra
+using Distributions
+import Mooncake
+import ForwardDiff
+
+using MeasureBase
+using MeasureBase: transport_to, transport_def, asmeasure
+using MeasureBase: StdUniform, StdNormal
+
+_mooncake_gradient(f, x) = Mooncake.value_and_gradient!!(
+    Mooncake.prepare_gradient_cache(f, x), f, x
+)[2][2]
+
+_test_gradient(f, x::Real) = @test _mooncake_gradient(f, x) ≈ ForwardDiff.derivative(f, x)
+_test_gradient(f, x::AbstractVector) = @test _mooncake_gradient(f, x) ≈ ForwardDiff.gradient(f, x)
+
+@testset "Mooncake AD with Distributions" begin
+    @test Base.get_extension(MeasureBase, :MeasureBaseDistributionsMooncakeExt) isa Module
+
+    @testset "zero-derivative primitives" begin
+        rng = Random.Xoshiro(789990641)
+        Mooncake.TestUtils.test_rule(
+            rng, MeasureBase._dist_params_numtype, Normal(0.2, 1.3);
+            is_primitive = true,
+        )
+    end
+
+    @testset "univariate transport gradients" begin
+        _test_gradient(x -> transport_def(StdUniform(), Normal(1.0, 2.0), x), 0.5)
+        _test_gradient(u -> transport_def(Normal(1.0, 2.0), StdUniform(), u), 0.3)
+        _test_gradient(u -> transport_def(Beta(2.0, 3.0), StdUniform(), u), 0.3)
+        _test_gradient(x -> transport_def(StdUniform(), Gamma(2.0, 1.0), x), 0.7)
+        _test_gradient(x -> transport_def(StdUniform(), truncated(Normal(0.3, 1.2), -0.5, 1.5), x), 0.4)
+        _test_gradient(x -> transport_def(StdUniform(), 2.0 * Weibull(0.7) + 1.0, x), 3.0)
+        _test_gradient(x -> transport_def(StdNormal(), StandardDist{Uniform}(), x), 0.4)
+    end
+
+    @testset "multivariate transport gradients" begin
+        mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3])
+        _test_gradient(x -> sum(transport_to(StdNormal()^2, mvn)(x)), [0.1, -2.0])
+        _test_gradient(y -> sum(transport_to(mvn, StdNormal()^2)(y)), [0.2, 0.7])
+
+        pd = product_distribution([Weibull(0.7), Exponential(1.3), Normal(0.5, 2.0)])
+        _test_gradient(x -> sum(transport_to(StdNormal()^3, asmeasure(pd))(x)), [0.4, 0.8, 1.5])
+
+        dirich = Dirichlet([2.0, 3.0, 4.0])
+        _test_gradient(u -> MeasureBase.from_origin(dirich, u)[1], [0.3, 0.7])
+        _test_gradient(x -> sum(MeasureBase.to_origin(dirich, vcat(x, 1 - sum(x)))), [0.28, 0.23])
+    end
+
+    @testset "logdensityof gradients" begin
+        for d in [
+            Weibull(0.7, 1.3),
+            MixtureModel([Normal(-1.0, 1.0), Normal(2.0, 3.0)], [0.3, 0.7]),
+            MixtureModel([Normal(-2.0, 1.0), Normal(0.0, 2.0), Normal(3.0, 1.0)], [0.2, 0.5, 0.3]),
+        ]
+            m = asmeasure(d)
+            _test_gradient(x -> logdensityof(m, x[1]), [0.5])
+        end
+
+        mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3])
+        _test_gradient(x -> logdensityof(asmeasure(mvn), x), [0.1, -2.0])
+    end
+end
diff --git a/test/distributions/test_standard_dist.jl b/test/distributions/test_standard_dist.jl
new file mode 100644
index 00000000..64b9f655
--- /dev/null
+++ b/test/distributions/test_standard_dist.jl
@@ -0,0 +1,128 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using Random, Statistics, LinearAlgebra
+using Distributions, PDMats
+using StableRNGs
+import ForwardDiff, ChainRulesTestUtils
+
+
+@testset "standard_dist" begin
+    stblrng() = StableRNG(789990641)
+
+    for (D, sz, dref) in [
+        (Uniform, (), Uniform()),
+        (Uniform, (5,), product_distribution(fill(Uniform(0.0, 1.0), 5))),
+        (Uniform, (2, 3), reshape(product_distribution(fill(Uniform(0.0, 1.0), 6)), 2, 3)),
+        (Normal, (), Normal()),
+        (Normal, (), Normal(0., 1.0)),
+        (Normal, (5,), MvNormal(Diagonal(fill(1.0, 5)))),
+        (Normal, (2, 3), reshape(MvNormal(Diagonal(fill(1.0, 6))), 2, 3)),
+        (Exponential, (), Exponential()),
+        (Exponential, (5,), product_distribution(fill(Exponential(1.0), 5))),
+        (Exponential, (2, 3), reshape(product_distribution(fill(Exponential(1.0), 6)), 2, 3)),
+    ]
+        @testset "StandardDist{$D}($(join(sz,",")))" begin
+            N = length(sz)
+
+            @test @inferred(StandardDist{D}(sz...)) isa StandardDist{D}
+            @test @inferred(StandardDist{D}(sz...)) isa StandardDist{D}
+            @test @inferred(size(StandardDist{D}(sz...))) == size(dref)
+            @test @inferred(size(StandardDist{D}(sz...))) == size(dref)
+
+            d = StandardDist{D}(sz...)
+
+            if size(d) == ()
+                @test @inferred(MeasureBaseDistributionsExt.nonstddist(d)) == dref
+            end
+
+            @test @inferred(length(d)) == length(dref)
+            @test @inferred(size(d)) == size(dref)
+
+            @test @inferred(eltype(typeof(d))) == eltype(typeof(dref))
+            @test @inferred(eltype(d)) == eltype(dref)
+
+            @test @inferred(Distributions.params(d)) == ()
+            @test @inferred(partype(d)) == partype(dref)
+
+            for f in [minimum, maximum, mean, median, mode, modes, var, std, skewness, kurtosis, location, scale, entropy]
+                supported_by_dref = try f(dref); true catch MethodError; false; end
+                if supported_by_dref
+                    @test @inferred(f(d)) ≈ f(dref)
+                end
+            end
+
+            for x in [rand(dref) for i in 1:10]
+                ref_gradlogpdf = try
+                    gradlogpdf(dref, x)
+                catch MethodError
+                    ForwardDiff.gradient(x -> logpdf(dref, x), x)
+                end
+                @test @inferred(gradlogpdf(d, x)) ≈ ref_gradlogpdf
+                @test @inferred(logpdf(d, x)) ≈ logpdf(dref, x)
+                @test @inferred(pdf(d, x)) ≈ pdf(dref, x)
+            end
+
+            if size(d) == ()
+                for x in [minimum(dref), quantile(dref, 1//3), quantile(dref, 1//2), quantile(dref, 2//3), maximum(dref)]
+                    for f in [logpdf, pdf, gradlogpdf, logcdf, cdf, logccdf, ccdf]
+                        @test @inferred(f(d, x)) ≈ f(dref, x)
+                    end
+                end
+
+                for x in [0, 1//3, 1//2, 2//3, 1]
+                    for f in [quantile, cquantile]
+                        @test @inferred(f(d, x)) ≈ f(dref, x)
+                    end
+                end
+
+                for x in log.([0, 1//3, 1//2, 2//3, 1])
+                    for f in [invlogcdf, invlogccdf]
+                        @test @inferred(f(d, x)) ≈ f(dref, x)
+                    end
+                end
+
+                for p in [0.0, 0.25, 0.75, 1.0]
+                    @test @inferred(quantile(d, p)) == quantile(dref, p)
+                    @test @inferred(cquantile(d, p)) == cquantile(dref, p)
+                end
+
+                for t in [-3, 0, 3]
+                    @test isapprox(@inferred(mgf(d, t)), mgf(dref, t), rtol = 1e-5)
+                    @test isapprox(@inferred(cf(d, t)), cf(dref, t), rtol = 1e-5)
+                end
+
+                @test @inferred(truncated(d, quantile(dref, 1//3), quantile(dref, 2//3))) == truncated(dref, quantile(dref, 1//3), quantile(dref, 2//3))
+
+                @test @inferred(product_distribution(fill(d, 3))) == StandardDist{typeof(d)}(3)
+                @test @inferred(product_distribution(fill(d, 3, 4))) == StandardDist{typeof(d)}(3, 4)
+            end
+
+            if length(size(d)) == 1
+                @test @inferred(convert(Distributions.Product, d)) isa Distributions.Product
+                d_as_prod = convert(Distributions.Product, d)
+                @test d_as_prod.v == fill(StandardDist{D}(), size(d)...)
+            end
+
+            @test @inferred(rand(stblrng(), d)) == rand(stblrng(), d)
+            @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), d, 5)
+
+            @test @inferred(rand(stblrng(), d)) == rand(stblrng(), dref)
+            @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), dref, 5)
+            @test @inferred(rand!(stblrng(), d, zeros(size(d)...))) == rand!(stblrng(), dref, zeros(size(dref)...))
+            if length(size(d)) == 1
+                @test @inferred(rand!(stblrng(), d, zeros(size(d)..., 5))) == rand!(stblrng(), dref, zeros(size(dref)..., 5))
+            end
+        end
+    end
+
+    @testset "StandardDist{Normal}()" begin
+        # TODO: Add @inferred
+        d = StandardDist{Normal}(4)
+        d_uv = StandardDist{Normal}()
+        dref = MvNormal(Diagonal(fill(1.0, 4)))
+        @test (MvNormal(d)) == dref
+        @test (Base.convert(MvNormal, d)) == dref
+    end
+end
diff --git a/test/distributions/test_standard_normal.jl b/test/distributions/test_standard_normal.jl
new file mode 100644
index 00000000..3d77f583
--- /dev/null
+++ b/test/distributions/test_standard_normal.jl
@@ -0,0 +1,129 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using Random, Statistics, LinearAlgebra
+using Distributions, PDMats
+using StableRNGs
+
+
+@testset "StandardDist{Normal}" begin
+    stblrng() = StableRNG(789990641)
+
+    @testset "StandardDist{Normal,0}" begin
+        @test @inferred(Normal(StandardDist{Normal}())) isa Normal{Float64}
+        @test @inferred(Normal(StandardDist{Normal}())) == Normal()
+        @test @inferred(convert(Normal, StandardDist{Normal}())) == Normal()
+
+        d = StandardDist{Normal}()
+        dref = Normal()
+
+        @test @inferred(minimum(d)) == minimum(dref)
+        @test @inferred(maximum(d)) == maximum(dref)
+        
+        @test @inferred(Distributions.params(d)) == ()
+        @test @inferred(partype(d)) == partype(dref)
+        
+        @test @inferred(location(d)) == location(dref)
+        @test @inferred(scale(d)) == scale(dref)
+        
+        @test @inferred(eltype(typeof(d))) == eltype(typeof(dref))
+        @test @inferred(eltype(d)) == eltype(dref)
+
+        @test @inferred(length(d)) == length(dref)
+        @test @inferred(size(d)) == size(dref)
+
+        @test @inferred(mean(d)) == mean(dref)
+        @test @inferred(median(d)) == median(dref)
+        @test @inferred(mode(d)) == mode(dref)
+        @test @inferred(modes(d)) ≈ modes(dref)
+        
+        @test @inferred(var(d)) == var(dref)
+        @test @inferred(std(d)) == std(dref)
+        @test @inferred(skewness(d)) == skewness(dref)
+        @test @inferred(kurtosis(d)) == kurtosis(dref)
+        
+        @test @inferred(entropy(d)) == entropy(dref)
+        
+        for x in [-Inf, -1.3, 0.0, 1.3, +Inf]
+            @test @inferred(gradlogpdf(d, x)) == gradlogpdf(dref, x)
+            
+            @test @inferred(logpdf(d, x)) == logpdf(dref, x)
+            @test @inferred(pdf(d, x)) == pdf(dref, x)
+            @test @inferred(logcdf(d, x)) == logcdf(dref, x)
+            @test @inferred(cdf(d, x)) == cdf(dref, x)
+            @test @inferred(logccdf(d, x)) == logccdf(dref, x)
+            @test @inferred(ccdf(d, x)) == ccdf(dref, x)
+        end
+
+        for p in [0.0, 0.25, 0.75, 1.0]
+            @test @inferred(quantile(d, p)) == quantile(dref, p)
+            @test @inferred(cquantile(d, p)) == cquantile(dref, p)
+        end
+
+        for t in [-3, 0, 3]
+            @test @inferred(mgf(d, t)) == mgf(dref, t)
+            @test @inferred(cf(d, t)) == cf(dref, t)
+        end
+
+        @test @inferred(rand(stblrng(), d)) == rand(stblrng(), dref)
+        @test @inferred(rand!(stblrng(), d, fill(0.0))) == rand!(stblrng(), dref, fill(0.0))
+        @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), dref, 5)
+
+        @test @inferred(truncated(StandardDist{Normal}(), -2.2f0, 3.1f0)) isa Truncated{Normal{Float64}}
+        @test truncated(StandardDist{Normal}(), -2.2f0, 3.1f0) == truncated(Normal(0.0, 1.0), -2.2f0, 3.1f0)
+
+        @test @inferred(product_distribution(fill(StandardDist{Normal}(), 3))) isa StandardDist{Normal,1}
+        @test product_distribution(fill(StandardDist{Normal}(), 3)) == StandardDist{Normal}(3)
+    end
+
+
+    @testset "StandardDist{Normal,1}" begin
+        @test @inferred(StandardDist{Normal}(3)) isa StandardDist{Normal,1}
+        @test @inferred(StandardDist{Normal}(3)) isa StandardDist{Normal,1}
+        @test @inferred(StandardDist{Normal}(3)) isa StandardDist{Normal,1}
+
+        @test @inferred(MvNormal(StandardDist{Normal}(3))) isa MvNormal{Int}
+        @test @inferred(MvNormal(StandardDist{Normal}(3))) == MvNormal(ScalMat(3, 1.0))
+        @test @inferred(convert(MvNormal, StandardDist{Normal}(3))) == MvNormal(ScalMat(3, 1.0))
+
+        d = StandardDist{Normal}(3)
+        dref = MvNormal(ScalMat(3, 1.0))
+
+        @test @inferred(eltype(typeof(d))) == eltype(typeof(dref))
+        @test @inferred(eltype(d)) == eltype(dref)
+
+        @test @inferred(length(d)) == length(dref)
+        @test @inferred(size(d)) == size(dref)
+
+        @test @inferred(Distributions.params(d)) == ()
+        @test @inferred(partype(d)) == partype(dref)
+
+        @test @inferred(mean(d)) == mean(dref)
+        @test @inferred(var(d)) == var(dref)
+        @test @inferred(cov(d)) == cov(dref)
+        
+        @test @inferred(mode(d)) == mode(dref)
+        @test @inferred(modes(d)) == modes(dref)
+
+        @test @inferred(invcov(d)) == invcov(dref)
+        @test @inferred(logdetcov(d)) == logdetcov(dref)
+
+        @test @inferred(entropy(d)) == entropy(dref)
+
+        for x in fill.([-Inf, -1.3, 0.0, 1.3, +Inf], 3)
+            # Distributions.insupport is inconsistent at +- Inf between Normal and MvNormal
+            if !any(isinf, x)
+                @test @inferred(Distributions.insupport(d, x)) == Distributions.insupport(dref, x)
+            end
+            @test @inferred(logpdf(d, x)) == logpdf(dref, x)
+            @test @inferred(pdf(d, x)) == pdf(dref, x)
+            @test @inferred(sqmahal(d, x)) == sqmahal(dref, x)
+            @test @inferred(gradlogpdf(d, x)) == gradlogpdf(dref, x)
+        end
+
+        @test @inferred(rand(stblrng(), d)) == rand(stblrng(), d)
+        @test @inferred(rand!(stblrng(), d, zeros(3))) == rand!(stblrng(), d, zeros(3))
+        @test @inferred(rand!(stblrng(), d, zeros(3, 10))) == rand!(stblrng(), d, zeros(3, 10))
+    end
+end
diff --git a/test/distributions/test_standard_uniform.jl b/test/distributions/test_standard_uniform.jl
new file mode 100644
index 00000000..bcb0fb3e
--- /dev/null
+++ b/test/distributions/test_standard_uniform.jl
@@ -0,0 +1,118 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using Random, Statistics, LinearAlgebra
+using Distributions, PDMats
+using StableRNGs
+using FillArrays
+using ForwardDiff
+
+
+@testset "StandardDist{Uniform}" begin
+    stblrng() = StableRNG(789990641)
+
+    @testset "StandardDist{Uniform,0}" begin
+        @test @inferred(Uniform(StandardDist{Uniform}())) isa Uniform{Float64}
+        @test @inferred(Uniform(StandardDist{Uniform}())) == Uniform()
+        @test @inferred(convert(Uniform, StandardDist{Uniform}())) == Uniform()
+
+        d = StandardDist{Uniform}()
+        dref = Uniform()
+
+        @test @inferred(minimum(d)) == minimum(dref)
+        @test @inferred(maximum(d)) == maximum(dref)
+        
+        @test @inferred(Distributions.params(d)) == ()
+        @test @inferred(partype(d)) == partype(dref)
+        
+        @test @inferred(location(d)) == location(dref)
+        @test @inferred(scale(d)) == scale(dref)
+        
+        @test @inferred(eltype(typeof(d))) == eltype(typeof(dref))
+        @test @inferred(eltype(d)) == eltype(dref)
+
+        @test @inferred(length(d)) == length(dref)
+        @test @inferred(size(d)) == size(dref)
+
+        @test @inferred(mean(d)) == mean(dref)
+        @test @inferred(median(d)) == median(dref)
+        @test @inferred(mode(d)) == mode(dref)
+        @test @inferred(modes(d)) ≈ modes(dref)
+        
+        @test @inferred(var(d)) ≈ var(dref)
+        @test @inferred(std(d)) ≈ std(dref)
+        @test @inferred(skewness(d)) == skewness(dref)
+        @test @inferred(kurtosis(d)) ≈ kurtosis(dref)
+        
+        @test @inferred(entropy(d)) == entropy(dref)
+        
+        for x in [-0.5, 0.0, 0.25, 0.75, 1.0, 1.5]
+            @test @inferred(logpdf(d, x)) == logpdf(dref, x)
+            @test @inferred(pdf(d, x)) == pdf(dref, x)
+            @test @inferred(logcdf(d, x)) == logcdf(dref, x)
+            @test @inferred(cdf(d, x)) == cdf(dref, x)
+            @test @inferred(logccdf(d, x)) == logccdf(dref, x)
+            @test @inferred(ccdf(d, x)) == ccdf(dref, x)
+        end
+
+        for p in [0.0, 0.25, 0.75, 1.0]
+            @test @inferred(quantile(d, p)) == quantile(dref, p)
+            @test @inferred(cquantile(d, p)) == cquantile(dref, p)
+        end
+
+        for t in [-3, 0, 3]
+            @test @inferred(mgf(d, t)) == mgf(dref, t)
+            @test @inferred(cf(d, t)) == cf(dref, t)
+        end
+
+        @test @inferred(rand(stblrng(), d)) == rand(stblrng(), dref)
+        @test @inferred(rand!(stblrng(), d, fill(0.0))) == rand!(stblrng(), dref, fill(0.0))
+        @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), dref, 5)
+
+        @test @inferred(truncated(StandardDist{Uniform}(), -0.5f0, 0.7f0)) isa Uniform{Float64}
+        @test truncated(StandardDist{Uniform}(), -0.5f0, 0.7f0) == Uniform(0.0f0, 0.7f0)
+        @test truncated(StandardDist{Uniform}(), 0.2f0, 0.7f0) == Uniform(0.2f0, 0.7f0)
+
+        @test @inferred(product_distribution(fill(StandardDist{Uniform}(), 3))) isa MeasureBaseDistributionsExt.StandardDist{Uniform,1}
+        @test product_distribution(fill(StandardDist{Uniform}(), 3)) == MeasureBaseDistributionsExt.StandardDist{Uniform}(3)
+    end
+
+
+    @testset "StandardDist{Uniform,1}" begin
+        d = MeasureBaseDistributionsExt.StandardDist{Uniform}(3)
+        dref = product_distribution(fill(Uniform(), 3))
+
+        @test @inferred(eltype(typeof(d))) == eltype(typeof(dref))
+        @test @inferred(eltype(d)) == eltype(dref)
+
+        @test @inferred(length(d)) == length(dref)
+        @test @inferred(size(d)) == size(dref)
+
+        @test @inferred(Distributions.params(d)) == ()
+        @test @inferred(partype(d)) == partype(dref)
+
+        @test @inferred(mean(d)) == mean(dref)
+        @test @inferred(var(d)) ≈ var(dref)
+        @test @inferred(cov(d)) ≈ cov(dref)
+
+        @test @inferred(mode(d)) == [0.5, 0.5, 0.5]
+        @test @inferred(modes(d)) == fill([0, 0,0 ])
+
+        @test @inferred(invcov(d)) == inv(cov(dref))
+        @test @inferred(logdetcov(d)) == logdet(cov(dref))
+
+        @test @inferred(entropy(d)) == entropy(dref)
+
+        for x in fill.([-Inf, -1.3, 0.0, 1.3, +Inf], 3)
+            @test @inferred(Distributions.insupport(d, x)) == Distributions.insupport(dref, x)
+            @test @inferred(logpdf(d, x)) == logpdf(dref, x)
+            @test @inferred(pdf(d, x)) == pdf(dref, x)
+            @test @inferred(gradlogpdf(d, x)) == ForwardDiff.gradient(x -> logpdf(d, x), x)
+        end
+
+        @test @inferred(rand(stblrng(), d)) == rand(stblrng(), d)
+        @test @inferred(rand!(stblrng(), d, zeros(3))) == rand!(stblrng(), d, zeros(3))
+        @test @inferred(rand!(stblrng(), d, zeros(3, 10))) == rand!(stblrng(), d, zeros(3, 10))
+    end
+end
diff --git a/test/distributions/test_transport.jl b/test/distributions/test_transport.jl
new file mode 100644
index 00000000..4ae50831
--- /dev/null
+++ b/test/distributions/test_transport.jl
@@ -0,0 +1,194 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using LinearAlgebra
+using InverseFunctions, ChangesOfVariables
+using Distributions, ArraysOfArrays
+using StableRNGs
+import ForwardDiff, Zygote
+import PDMats
+
+using MeasureBase: transport_to, transport_def, transport_origin
+using MeasureBase: StdUniform, StdNormal, StdExponential
+using .MeasureBaseDistributionsExt: _trafo_cdf, _trafo_quantile
+
+include("getjacobian.jl")
+
+
+@testset "test_distribution_transform" begin
+    function test_back_and_forth(trg, src)
+        @testset "transform $(typeof(trg).name) <-> $(typeof(src).name)" begin
+            x = rand(src)
+            y = transport_def(trg, src, x)
+            src_v_reco = transport_def(src, trg, y)
+
+            @test x ≈ src_v_reco
+            
+            f = x -> transport_def(trg, src, x)
+            ref_ladj = logpdf(src, x) - logpdf(trg, y)
+            @test ref_ladj ≈ logabsdet(getjacobian(f, x))[1]
+        end
+    end
+
+    reshaped_rand(d::Distribution{Univariate}, n) = rand(d, n)
+    reshaped_rand(d::Distribution{Multivariate}, n) = nestedview(rand(d, n))
+
+    function test_dist_trafo_moments(trg, src)
+        unshaped(x) = first(torv_and_back(x))
+        @testset "check moments of trafo $(typeof(trg).name) <- $(typeof(src).name)" begin
+            X = reshaped_rand(src, 10^5)
+            Y = transport_to(trg, src).(X)
+            Y_ref = reshaped_rand(trg, 10^6)
+            @test isapprox(mean(unshaped.(Y)), mean(unshaped.(Y_ref)), rtol = 0.5)
+            @test isapprox(cov(unshaped.(Y)), cov(unshaped.(Y_ref)), rtol = 0.5)
+        end
+    end
+
+    @testset "transforms-tests" begin
+        stduvuni = StandardDist{Uniform}()
+        stduvnorm = StandardDist{Uniform}()
+
+        uniform1 = Uniform(-5.0, -0.01)
+        uniform2 = Uniform(0.01, 5.0)
+
+        normal1 = Normal(-10, 1)
+        normal2 = Normal(10, 5)
+
+        stdmvnorm1 = StandardDist{Normal}(1)
+        stdmvnorm2 = StandardDist{Normal}(2)
+
+        stdmvuni2 = StandardDist{Uniform}(2)
+
+        standnorm2_reshaped = reshape(stdmvnorm2, 1, 2)
+
+        mvnorm = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3])
+        beta = Beta(3,1)
+        gamma = Gamma(0.1,0.7)
+        dirich = Dirichlet([0.1,4])
+
+        test_back_and_forth(stduvuni, stduvuni)
+        test_back_and_forth(stduvnorm, stduvnorm)
+        test_back_and_forth(stduvuni, stduvnorm)
+        test_back_and_forth(stduvnorm, stduvuni)
+
+        test_back_and_forth(stdmvuni2, stdmvuni2)
+        test_back_and_forth(stdmvnorm2, stdmvnorm2)
+        test_back_and_forth(stdmvuni2, stdmvnorm2)
+        test_back_and_forth(stdmvnorm2, stdmvuni2)
+
+        test_back_and_forth(beta, stduvnorm)
+        test_back_and_forth(gamma, stduvnorm)
+        test_back_and_forth(gamma, beta)
+
+        test_back_and_forth(mvnorm, stdmvuni2)
+        test_back_and_forth(stdmvuni2, mvnorm)
+
+        test_back_and_forth(mvnorm, standnorm2_reshaped)
+        test_back_and_forth(standnorm2_reshaped, mvnorm)
+        test_back_and_forth(stdmvnorm2, standnorm2_reshaped)
+        test_back_and_forth(standnorm2_reshaped, standnorm2_reshaped)
+
+        test_dist_trafo_moments(normal2, normal1)
+        test_dist_trafo_moments(uniform2, uniform1)
+
+        test_dist_trafo_moments(beta, stduvnorm)
+        test_dist_trafo_moments(gamma, stduvnorm)
+
+        test_dist_trafo_moments(mvnorm, stdmvnorm2)
+        test_dist_trafo_moments(dirich, stdmvnorm1)
+
+        let
+            mvuni = product_distribution([Uniform(), Uniform()])
+
+            x = rand()
+            @test_throws ArgumentError transport_to(stduvnorm, mvnorm)(x)
+            @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm1)(x)
+            @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm2)(x)
+
+            x = rand(2)
+            @test_throws ArgumentError transport_to(stduvnorm, mvnorm)(x)
+            @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm1)(x)
+            @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm2)(x)
+        end
+    end
+
+    @testset "Custom cdf and quantile for dual numbers" begin
+        Dual = ForwardDiff.Dual
+
+        @test isapprox(_trafo_cdf(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), cdf(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), rtol = 10^-6)
+        @test isapprox(_trafo_cdf(Normal(0, 1), Dual(0.5, 1)), cdf(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6)
+
+        @test isapprox(_trafo_quantile(Normal(0, 1), Dual(0.5, 1)), quantile(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6)
+        @test isapprox(_trafo_quantile(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), quantile(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), rtol = 10^-6)
+    end
+
+    @testset "trafo autodiff pullbacks" begin
+        x = [0.6, 0.7, 0.8, 0.9]
+        f = transport_to(Dirichlet([3.0, 4.0, 5.0, 6.0, 7.0]), Uniform)
+        @test isapprox(ForwardDiff.jacobian(f, x), Zygote.jacobian(f, x)[1], rtol = 10^-4)
+        f = inverse(transport_to(Normal, Dirichlet([3.0, 4.0, 5.0, 6.0, 7.0])))
+        @test isapprox(ForwardDiff.jacobian(f, x), Zygote.jacobian(f, x)[1], rtol = 10^-4)
+    end
+
+
+    @testset "transport_to autosel" begin
+        for (M,R) in [
+            (StandardDist{Normal}, StandardDist{Normal})
+            (Normal, StandardDist{Normal})
+            (StandardDist{Uniform}, StandardDist{Uniform})
+            (Uniform, StandardDist{Uniform})
+        ]
+            @test @inferred(transport_to(M, Weibull())) == transport_to(R(), Weibull())
+            @test @inferred(transport_to(Weibull(), M)) == transport_to(Weibull(), R())
+            @test @inferred(transport_to(M, MvNormal(float(I(5))))) == transport_to(R(5), MvNormal(float(I(5))))
+            @test @inferred(transport_to(MvNormal(float(I(5))), M)) == transport_to(MvNormal(float(I(5))), R(5))
+            @test @inferred(transport_to(M, StdExponential()^(2,3))) == transport_to(R(6), StdExponential()^(2,3))
+            @test @inferred(transport_to(StdExponential()^(2,3), M)) == transport_to(StdExponential()^(2,3), R(6))
+        end
+    end
+
+    @testset "affine transformed distributions" begin
+        d = 2.0 * Weibull(0.7) + 1.0
+        x = rand(StableRNG(789990641), d)
+        u = transport_to(StdUniform(), d)(x)
+        @test u ≈ cdf(d, x)
+        @test transport_to(d, StdUniform())(u) ≈ x
+        test_back_and_forth(StandardDist{Normal}(), d)
+    end
+
+    @testset "truncated distributions" begin
+        d = truncated(Normal(0.3, 1.2), -0.5, 1.5)
+        for u in [0.0, 0.25, 0.75, 1.0, prevfloat(1.0)]
+            x = transport_to(d, StdUniform())(u)
+            @test minimum(d) <= x <= maximum(d)
+        end
+        test_back_and_forth(StandardDist{Uniform}(), d)
+    end
+
+    @testset "products of distributions" begin
+        pd = product_distribution([Weibull(0.7), Exponential(1.3), Normal(0.5, 2.0)])
+        m = MeasureBase.asmeasure(pd)
+        x = rand(StableRNG(789990641), pd)
+        for trg in [StdUniform()^3, StdNormal()^3]
+            y = transport_to(trg, m)(x)
+            y_ref = map((d_i, x_i) -> transport_to(trg.parent, d_i)(x_i), pd.v, x)
+            @test y ≈ y_ref
+            @test transport_to(m, trg)(y) ≈ x
+        end
+
+        pd2 = product_distribution([Normal(2.0, 0.5), Weibull(1.2), Uniform(-1.0, 3.0)])
+        m2 = MeasureBase.asmeasure(pd2)
+        y = transport_to(m2, m)(x)
+        @test transport_to(m, m2)(y) ≈ x
+    end
+
+    @testset "MvNormal covariance representations" begin
+        for Σ in [PDMats.ScalMat(3, 2.5), PDMats.PDiagMat([0.5, 1.0, 2.5]), Diagonal([0.5, 1.0, 2.5])]
+            mvn = MvNormal([0.2, -0.4, 0.6], Σ)
+            x = rand(StableRNG(789990641), mvn)
+            y = transport_to(StandardDist{Normal}(3), mvn)(x)
+            @test transport_to(mvn, StandardDist{Normal}(3))(y) ≈ x
+        end
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index ff5c3140..f0d4488a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -9,7 +9,6 @@ using MeasureBase: test_interface, test_smf
 
 include("test_aqua.jl")
 
-include("static.jl")
 
 include("test_primitive.jl")
 include("test_standard.jl")
@@ -26,4 +25,6 @@ include("combinators/transformedmeasure.jl")
 include("combinators/reshape.jl")
 include("combinators/implicitlymapped.jl")
 
+include("distributions/test_distributions.jl")
+
 include("test_docs.jl")
diff --git a/test/test_aqua.jl b/test/test_aqua.jl
index b6290e31..d4546ac2 100644
--- a/test/test_aqua.jl
+++ b/test/test_aqua.jl
@@ -9,5 +9,10 @@ import MeasureBase
 #end # testset
 
 Test.@testset "Aqua tests" begin
-    Aqua.test_all(MeasureBase, ambiguities = false)
+    Aqua.test_all(
+        MeasureBase,
+        ambiguities = false,
+        # Only used by package extensions:
+        stale_deps = (ignore = [:ArgCheck, :ArraysOfArrays],),
+    )
 end # testset

From 6f51ef626317e293d59b98eb52b6d38c4149bc18 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 1 Nov 2024 20:28:07 +0100
Subject: [PATCH 026/122] Remove PowerWeightedMeasure

Unused and untested.

(cherry picked from commit 1d271885dfeeac1149817af1b748588b67a4d696)
---
 src/MeasureBase.jl               |  1 -
 src/combinators/powerweighted.jl | 37 --------------------------------
 2 files changed, 38 deletions(-)
 delete mode 100644 src/combinators/powerweighted.jl

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 4ab59601..e3e4cf5e 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -194,7 +194,6 @@ include("combinators/likelihood.jl")
 include("combinators/pointwise.jl")
 include("combinators/restricted.jl")
 include("combinators/smart-constructors.jl")
-include("combinators/powerweighted.jl")
 include("combinators/conditional.jl")
 include("combinators/implicitlymapped.jl")
 
diff --git a/src/combinators/powerweighted.jl b/src/combinators/powerweighted.jl
deleted file mode 100644
index 47f50da4..00000000
--- a/src/combinators/powerweighted.jl
+++ /dev/null
@@ -1,37 +0,0 @@
-export ↑
-
-struct PowerWeightedMeasure{M,A} <: AbstractMeasure
-    parent::M
-    exponent::A
-end
-
-logdensity_def(d::PowerWeightedMeasure, x) = d.exponent * logdensity_def(d.parent, x)
-
-basemeasure(d::PowerWeightedMeasure, x) = basemeasure(d.parent, x)↑d.exponent
-
-basemeasure(d::PowerWeightedMeasure) = basemeasure(d.parent)↑d.exponent
-
-function powerweightedmeasure(d, α)
-    isone(α) && return d
-    PowerWeightedMeasure(d, α)
-end
-
-(d::AbstractMeasure)↑α = powerweightedmeasure(d, α)
-
-insupport(d::PowerWeightedMeasure, x) = insupport(d.parent, x)
-
-function Base.show(io::IO, d::PowerWeightedMeasure)
-    print(io, d.parent, " ↑ ", d.exponent)
-end
-
-function powerweightedmeasure(d::PowerWeightedMeasure, α)
-    powerweightedmeasure(d.parent, α * d.exponent)
-end
-
-function powerweightedmeasure(d::WeightedMeasure, α)
-    weightedmeasure(α * d.logweight, powerweightedmeasure(d.base, α))
-end
-
-function Pretty.tile(d::PowerWeightedMeasure)
-    Pretty.pair_layout(Pretty.tile(d.parent), Pretty.tile(d.exponent), sep = " ↑ ")
-end

From eac759b51282493f94b2f449c93b86c19fc8ab73 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 1 Nov 2024 20:28:07 +0100
Subject: [PATCH 027/122] Remove kernelfactor

Not used currently.

(cherry picked from commit 98445838a025ddfc310e59f0e4470baf6263f6b1)
---
 src/parameterized.jl | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/src/parameterized.jl b/src/parameterized.jl
index 78e43995..8b1c8c88 100644
--- a/src/parameterized.jl
+++ b/src/parameterized.jl
@@ -127,14 +127,3 @@ params(::Type{PM}) where {N,PM<:ParameterizedMeasure{N}} = N
 function paramnames(μ, constraints::NamedTuple{N}) where {N}
     tuple((k for k in paramnames(μ) if k ∉ N)...)
 end
-
-###############################################################################
-# kernelfactor
-
-function kernelfactor(::Type{P}) where {N,P<:ParameterizedMeasure{N}}
-    (constructorof(P), N)
-end
-
-function kernelfactor(::P) where {N,P<:ParameterizedMeasure{N}}
-    (constructorof(P), N)
-end

From dc94b4a474e2312257a40f8e6c5438c6b65bccf1 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 1 Nov 2024 20:28:08 +0100
Subject: [PATCH 028/122] Remove the rebase function

A rebase can easily be written explicitly.

(cherry picked from commit fb3c98cecb995e48b64367ccc7122081c9113b2b)
---
 src/MeasureBase.jl |  1 -
 src/density.jl     | 11 -----------
 2 files changed, 12 deletions(-)

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index e3e4cf5e..ecfbab2b 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -61,7 +61,6 @@ import HeterogeneousComputing
 using HeterogeneousComputing: real_numtype
 
 export gentype
-export rebase
 
 export AbstractMeasure
 
diff --git a/src/density.jl b/src/density.jl
index a79021de..35147645 100644
--- a/src/density.jl
+++ b/src/density.jl
@@ -180,14 +180,3 @@ function logdensityof(μ::DensityMeasure, x::Any)
         convert(R, integrand_logval + base_logval)::R
     end
 end
-
-"""
-    rebase(μ, ν)
-
-Express `μ` in terms of a density over `ν`. Satisfies
-```
-basemeasure(rebase(μ, ν)) == ν
-density(rebase(μ, ν)) == 𝒹(μ,ν)
-``` 
-"""
-rebase(μ, ν) = ∫(𝒹(μ, ν), ν)

From bb9c0213a065f9b790d47d305d1595370940c34b Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 1 Nov 2024 20:28:08 +0100
Subject: [PATCH 029/122] Removes PointwiseProductMeasure

`mintegral` should be used instead to express posteriors.

(cherry picked from commit 3c611806cd54740d2458a5192c97f77085c186cb)
---
 src/MeasureBase.jl           |  1 -
 src/combinators/pointwise.jl | 30 ------------------------------
 2 files changed, 31 deletions(-)
 delete mode 100644 src/combinators/pointwise.jl

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index ecfbab2b..405114aa 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -190,7 +190,6 @@ include("combinators/product.jl")
 include("combinators/power.jl")
 include("combinators/spikemixture.jl")
 include("combinators/likelihood.jl")
-include("combinators/pointwise.jl")
 include("combinators/restricted.jl")
 include("combinators/smart-constructors.jl")
 include("combinators/conditional.jl")
diff --git a/src/combinators/pointwise.jl b/src/combinators/pointwise.jl
deleted file mode 100644
index 778e7f4e..00000000
--- a/src/combinators/pointwise.jl
+++ /dev/null
@@ -1,30 +0,0 @@
-export ⊙
-
-struct PointwiseProductMeasure{P,L} <: AbstractMeasure
-    prior::P
-    likelihood::L
-end
-
-iterate(p::PointwiseProductMeasure, i = 1) = iterate((p.prior, p.likelihood), i)
-
-function Pretty.tile(d::PointwiseProductMeasure)
-    Pretty.pair_layout(Pretty.tile(d.prior), Pretty.tile(d.likelihood), sep = " ⊙ ")
-end
-
-⊙(prior, ℓ) = pointwiseproduct(prior, ℓ)
-
-@inbounds function insupport(d::PointwiseProductMeasure, p)
-    prior, ℓ = d
-    istrue(insupport(prior, p)) && istrue(insupport(ℓ, p))
-end
-
-@inline function logdensity_def(d::PointwiseProductMeasure, p)
-    prior, ℓ = d
-    unsafe_logdensityof(ℓ, p)
-end
-
-basemeasure(d::PointwiseProductMeasure) = d.prior
-
-function gentype(d::PointwiseProductMeasure)
-    gentype(d.prior)
-end

From 48f5e934113816021a3c8a52e42d2263a1fed299 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 1 Nov 2024 20:28:08 +0100
Subject: [PATCH 030/122] Remove operator otimes

To be re-introduced in sub-module MeasureOperators.

(cherry picked from commit 0cdca3d443b14b28313264f17bfdc085a190c394)
---
 src/combinators/product.jl | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 9135dc2b..dbdfffa3 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -179,18 +179,6 @@ function testvalue(::Type{T}, d::AbstractProductMeasure) where {T}
     _map(m -> testvalue(T, m), marginals(d))
 end
 
-export ⊗
-
-"""
-    ⊗(μs::AbstractMeasure...)
-
-`⊗` is a binary operator for building product measures. This satisfies the law
-
-```
-    basemeasure(μ ⊗ ν) == basemeasure(μ) ⊗ basemeasure(ν)
-```
-"""
-⊗(μs::AbstractMeasure...) = productmeasure(μs)
 
 ###############################################################################
 # I <: Base.Generator

From f8e1a509eb05a12e50e50c9c54c0c7d93fc86cdc Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 1 Nov 2024 20:28:08 +0100
Subject: [PATCH 031/122] Remove scrd operator

To be reintroduced in submodule MeasureOperators

(cherry picked from commit 867adbeeeb0fe2bad0a2d691ddb54b050d0aa8e5)
---
 src/density.jl | 23 +----------------------
 1 file changed, 1 insertion(+), 22 deletions(-)

diff --git a/src/density.jl b/src/density.jl
index 35147645..0db02c29 100644
--- a/src/density.jl
+++ b/src/density.jl
@@ -20,8 +20,7 @@ For measures `μ` and `ν`, `Density(μ,ν)` represents the _density function_
 `dμ/dν`, also called the _Radon-Nikodym derivative_:
 https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem#Radon%E2%80%93Nikodym_derivative
 
-Instead of calling this directly, users should call `density_rel(μ, ν)` or
-its abbreviated form, `𝒹(μ,ν)`.
+Instead of calling this directly, users should call `density_rel(μ, ν)`.
 """
 struct Density{M,B} <: AbstractDensity
     μ::M
@@ -32,16 +31,6 @@ Base.:∘(::typeof(log), d::Density) = logdensity_rel(d.μ, d.base)
 
 Base.log(d::Density) = log ∘ d
 
-export 𝒹
-
-"""
-    𝒹(μ, base)
-
-Compute the density (Radon-Nikodym derivative) of μ with respect to `base`. This
-is a shorthand form for `density_rel(μ, base)`.
-"""
-𝒹(μ, base) = density_rel(μ, base)
-
 density_rel(μ, base) = Density(μ, base)
 
 (f::Density)(x) = density_rel(f.μ, f.base, x)
@@ -73,16 +62,6 @@ Base.:∘(::typeof(exp), d::LogDensity) = density_rel(d.μ, d.base)
 
 Base.exp(d::LogDensity) = exp ∘ d
 
-export log𝒹
-
-"""
-    log𝒹(μ, base)
-
-Compute the log-density (Radon-Nikodym derivative) of μ with respect to `base`.
-This is a shorthand form for `logdensity_rel(μ, base)`
-"""
-log𝒹(μ, base) = logdensity_rel(μ, base)
-
 logdensity_rel(μ, base) = LogDensity(μ, base)
 
 (f::LogDensity)(x) = logdensity_rel(f.μ, f.base, x)

From 7f32b7e66d2b9cc7d911346e12f94b1961040224 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 11:00:59 +0200
Subject: [PATCH 032/122] Rename bind to mbind and remove fish operator
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Combines the original commits f500aea, 86db05e and 58f0746 from the
measure-algebra branch: `mbind` takes the kernel as the first argument
(do-block support), the `↣` operator is removed (it looks very similar
to the `>=>` "fish" operator, which is not a monadic bind) and `Bind`
stores the kernel first.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/combinators/bind.jl | 53 +++++++++++++++++++++++------------------
 1 file changed, 30 insertions(+), 23 deletions(-)

diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl
index cc2022f2..60b7cfd8 100644
--- a/src/combinators/bind.jl
+++ b/src/combinators/bind.jl
@@ -1,36 +1,43 @@
+"""
+    struct MeasureBase.Bind{M,K} <: AbstractMeasure
+
+Represents a monatic bind. User code should not create instances of `Bind`
+directly, but should call `mbind(k, μ)` instead.
+"""
 struct Bind{M,K} <: AbstractMeasure
-    μ::M
     k::K
+    μ::M
+end
+
+getdof(d::Bind) = NoDOF{typeof(d)}()
+
+function Base.rand(rng::AbstractRNG, ::Type{T}, d::Bind) where {T}
+    x = rand(rng, T, d.μ)
+    y = rand(rng, T, d.k(x))
+    return y
 end
 
-export ↣
 
 """
-If 
-- μ is an `AbstractMeasure` or satisfies the Measure interface, and
-- k is a function taking values from the support of μ and returning a measure
+    mbind(k, μ)::AbstractMeasure
 
-Then `μ ↣ k` is a measure, called a *monadic bind*. In a
-probabilistic programming language like Soss.jl, this could be expressed as
+Given
 
-Note that bind is usually written `>>=`, but this symbol is unavailable in Julia.
+- a measure μ
+- a kernel function k that takes values from the support of μ and returns a
+  measure
+
+The *monadic bind* operation `mbind(k, μ)` returns is a new measure.
+
+A monadic bind is often written as `>>=` (e.g. in Haskell), but this symbol is
+unavailable in Julia.
 
 ```
-bind = @model μ,k begin 
-    x ~ μ
-    y ~ k(x)
-    return y
+μ = StdExponential()
+ν = mbind(μ) do scale
+    pushfwd(Base.Fix1(*, scale), StdNormal())
 end
 ```
-
-See also `bind` and `Bind`
 """
-↣(μ, k) = bind(μ, k)
-
-bind(μ, k) = Bind(μ, k)
-
-function Base.rand(rng::AbstractRNG, ::Type{T}, d::Bind) where {T}
-    x = rand(rng, T, d.μ)
-    y = rand(rng, T, d.k(x))
-    return y
-end
+mbind(k, μ) = Bind(k, μ)
+export mbind

From 81b0e7993b0fd9fc27e02b3cc95df0e155508655 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 11:04:12 +0200
Subject: [PATCH 033/122] Introduce mintegrate and mintegrate_exp
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Replaces the operators ∫ and ∫exp (which will return in the new
MeasureOperators submodule) with the functions `mintegrate` and
`mintegrate_exp`, and introduces `MeasureBase.as_integrand` and
`MeasureBase.as_integrand_exp` as specializable integrand conversion
hooks (folded forward from the later Likelihood/mintegrate refactor on
the measure-algebra branch). Adapts the Distributions extension and the
tests to use the function names.

Combines the original measure-algebra commits fefe36d, 74df0a5 and the
mintegrate-related parts of 1ed3f95.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../measure_interface.jl                      |   3 +-
 src/density.jl                                | 132 ++++++++++++++----
 test/distributions/test_measure_interface.jl  |   2 +-
 test/test_basics.jl                           |  16 +--
 4 files changed, 117 insertions(+), 36 deletions(-)

diff --git a/ext/MeasureBaseDistributionsExt/measure_interface.jl b/ext/MeasureBaseDistributionsExt/measure_interface.jl
index 6fed5d4a..933a3f6f 100644
--- a/ext/MeasureBaseDistributionsExt/measure_interface.jl
+++ b/ext/MeasureBaseDistributionsExt/measure_interface.jl
@@ -23,4 +23,5 @@
     Counting(MeasureBase.BoundedInts(static(0), static(Inf)))^size(d)
 
 
-MeasureBase.∫(f, base::Distribution) = MeasureBase.∫(f, convert(AbstractMeasure, base))
+MeasureBase.mintegrate(f, base::Distribution) =
+    MeasureBase.mintegrate(f, convert(AbstractMeasure, base))
diff --git a/src/density.jl b/src/density.jl
index 0db02c29..06dc98e1 100644
--- a/src/density.jl
+++ b/src/density.jl
@@ -50,8 +50,7 @@ For measures `μ` and `ν`, `LogDensity(μ,ν)` represents the _log-density func
 `log(dμ/dν)`, also called the _Radon-Nikodym derivative_:
 https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem#Radon%E2%80%93Nikodym_derivative
 
-Instead of calling this directly, users should call `logdensity_rel(μ, ν)` or
-its abbreviated form, `log𝒹(μ,ν)`.
+Instead of calling this directly, users should call `logdensity_rel(μ, ν)`.
 """
 struct LogDensity{M,B} <: AbstractDensity
     μ::M
@@ -77,12 +76,13 @@ DensityInterface.funcdensity(d::LogDensity) = throw(MethodError(funcdensity, (d,
         base    :: B
     end
 
-A `DensityMeasure` is a measure defined by a density or log-density with respect
-to some other "base" measure.
+A `DensityMeasure` is a measure defined by a density or log-density with
+respect to some other "base" measure.
 
-Users should not call `DensityMeasure` directly, but should instead call `∫(f,
-base)` (if `f` is a density function or `DensityInterface.IsDensity` object) or
-`∫exp(f, base)` (if `f` is a log-density function).
+Users should not instantiate `DensityMeasure` directly, but should instead
+call `mintegrate(f, base)` (if `f` is a density function or
+`DensityInterface.IsDensity` object) or `mintegrate_exp(f, base)` (if `f`
+is a log-density function).
 """
 struct DensityMeasure{F,B} <: AbstractMeasure
     f::F
@@ -99,42 +99,124 @@ end
 end
 
 function Pretty.tile(μ::DensityMeasure{F,B}) where {F,B}
-    result = Pretty.literal("DensityMeasure ∫(")
+    result = Pretty.literal("mintegrate(")
     result *= Pretty.pair_layout(Pretty.tile(μ.f), Pretty.tile(μ.base); sep = ", ")
     result *= Pretty.literal(")")
 end
 
-export ∫
 
 """
-    ∫(f, base::AbstractMeasure)
+    MeasureBase.as_integrand(f)
+    MeasureBase.as_integrand(density)
 
-Define a new measure in terms of a density `f` over some measure `base`.
+Make `f` or `density` (more) suitable as an integrand for
+[`mintegrate`](@ref).
+
+`mintegrate(obj, μ::AbstractMeasure)` automatically calls
+`as_integrand(obj)` internally.
+
+If a density is passed, it must implement the DensityInterface API.
+
+By default just returns `f` resp. `density`, but may be specialized for
+functions and densities that can profit from conversion to a form optimized
+for use in `mintegrate`.
+
+See also [`MeasureBase.as_likelihood`](@ref).
 """
-∫(f, base) = _densitymeasure(f, base, DensityKind(f))
+function as_integrand end
+
+@inline as_integrand(obj) = _as_integrand_default_impl(obj, DensityKind(obj))
+
+@inline _as_integrand_default_impl(f, ::NoDensity) = funcdensity(f)
+
+@inline _as_integrand_default_impl(density, ::IsDensity) = density
 
-_densitymeasure(f, base, ::IsDensity) = DensityMeasure(f, base)
-function _densitymeasure(f, base, ::HasDensity)
-    @error "`∫(f, base)` requires `DensityKind(f)` to be `IsDensity()` or `NoDensity()`."
+function _as_integrand_default_impl(obj, ::HasDensity)
+    throw(
+        ArgumentError(
+            "`MeasureBase.as_integrand(obj)` requires `DensityKind(obj)` to be `IsDensity()` or `NoDensity()`.",
+        ),
+    )
 end
-_densitymeasure(f, base, ::NoDensity) = DensityMeasure(funcdensity(f), base)
 
-export ∫exp
+
+@doc raw"""
+    mintegrate(f, μ::AbstractMeasure)::AbstractMeasure
+    mintegrate(density, μ::AbstractMeasure)::AbstractMeasure
+
+Returns a new measure that represents the indefinite
+[integral](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem)
+of `f` with respect to `μ`.
+
+If a density is passed, it must implement the DensityInterface API.
+
+`ν = mintegrate(f, μ)` generates a measure `ν` that has the mathematical
+interpretation
+
+```math
+\nu(A) = \int_A f(a) \, \rm{d}\mu(a)
+```
+"""
+function mintegrate end
+export mintegrate
+
+@inline mintegrate(obj, μ::AbstractMeasure) = DensityMeasure(as_integrand(obj), μ)
+
 
 """
-    ∫exp(f, base::AbstractMeasure)
+    MeasureBase.as_integrand_exp(log_f)
+
+Convert the logarithm of an integrand to an integrand.
 
-Define a new measure in terms of a log-density `f` over some measure `base`.
+See also [`MeasureBase.as_integrand`](@ref).
 """
-∫exp(f, base) = _logdensitymeasure(f, base, DensityKind(f))
+function as_integrand_exp end
+
+@inline as_integrand_exp(log_f) = _as_integrand_exp_default_impl(log_f, DensityKind(log_f))
+
+@inline _as_integrand_exp_default_impl(log_f, ::NoDensity) = logfuncdensity(log_f)
 
-function _logdensitymeasure(f, base, ::IsDensity)
-    @error "`∫exp(f, base)` is not valid when `DensityKind(f) == IsDensity()`. Use `∫(f, base)` instead."
+function _as_integrand_exp_default_impl(log_f, ::IsDensity)
+    throw(
+        ArgumentError(
+            "`as_integrand_exp(log_f)` is not valid when `DensityKind(log_f) == IsDensity()`. Use `as_integrand(log_f)` instead.",
+        ),
+    )
 end
-function _logdensitymeasure(f, base, ::HasDensity)
-    @error "`∫exp(f, base)` is not valid when `DensityKind(f) == HasDensity()`."
+
+function _as_integrand_exp_default_impl(log_f, ::HasDensity)
+    throw(
+        ArgumentError(
+            "`as_integrand_exp(log_f)` is not valid when `DensityKind(log_f) == HasDensity()`.",
+        ),
+    )
 end
-_logdensitymeasure(f, base, ::NoDensity) = DensityMeasure(logfuncdensity(f), base)
+
+
+@doc raw"""
+    mintegrate_exp(log_f, μ::AbstractMeasure)
+
+Given a function `log_f` that semantically represents the log of a function
+`f`, `mintegrate_exp` returns a new measure that represents the indefinite
+[integral](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem)
+of `f` with respect to `μ`.
+
+`ν = mintegrate_exp(log_f, μ)` generates a measure `ν` that has the
+mathematical interpretation
+
+```math
+\nu(A) = \int_A e^{log(f(a))} \, \rm{d}\mu(a) = \int_A f(a) \, \rm{d}\mu(a)
+```
+
+Note that `exp(log_f(...))` is usually not run explicitly, calculations that
+involve the resulting measure are typically performed in log-space,
+internally.
+"""
+function mintegrate_exp end
+export mintegrate_exp
+
+mintegrate_exp(log_f, μ::AbstractMeasure) = DensityMeasure(as_integrand_exp(log_f), μ)
+
 
 basemeasure(μ::DensityMeasure) = μ.base
 
diff --git a/test/distributions/test_measure_interface.jl b/test/distributions/test_measure_interface.jl
index d11d2889..f3d6c237 100644
--- a/test/distributions/test_measure_interface.jl
+++ b/test/distributions/test_measure_interface.jl
@@ -39,5 +39,5 @@ import MeasureBase
         @test (MeasureBase.params(c0)) == (α = 0.7, θ = 1.3)
     end
 
-    @test MeasureBase.∫(x -> Distributions.Normal(x, 0), Distributions.Normal()) isa MeasureBase.DensityMeasure
+    @test MeasureBase.mintegrate(x -> Distributions.Normal(x, 0), Distributions.Normal()) isa MeasureBase.DensityMeasure
 end
diff --git a/test/test_basics.jl b/test/test_basics.jl
index bd5a409c..11f1a8fe 100644
--- a/test/test_basics.jl
+++ b/test/test_basics.jl
@@ -1,4 +1,4 @@
-d = ∫exp(x -> -x^2, Lebesgue(ℝ))
+d = mintegrate_exp(x -> -x^2, Lebesgue(ℝ))
 
 # function draw2(μ)
 #     x = rand(μ)
@@ -125,7 +125,7 @@ end
           logdensityof(Lebesgue()^(3, 1), fill(2, 3, 1))
 end
 
-NormalMeasure() = ∫exp(x -> -0.5x^2, Lebesgue(ℝ))
+NormalMeasure() = mintegrate_exp(x -> -0.5x^2, Lebesgue(ℝ))
 
 @testset "Half" begin
     HalfNormal() = Half(NormalMeasure())
@@ -136,12 +136,10 @@ end
 
 @testset "Likelihood" begin
     ℓ = Likelihood(3) do (μ,)
-        ∫exp(Lebesgue(ℝ)) do x
+        mintegrate_exp(Lebesgue(ℝ)) do x
             -(x - μ)^2
         end
     end
-
-    @inferred logdensityof(Lebesgue() ⊙ ℓ, 2.0)
 end
 
 # @testset "Likelihood" begin
@@ -197,7 +195,7 @@ end
     f2 = x -> sqrt(abs(sum(x)))
     f3 = x -> 2 * sum(x)
     f4 = x -> sum(sqrt.(abs.(x)))
-    m = @inferred ∫exp(f1, ∫exp(f2, ∫exp(f3, ∫exp(f4, StdUniform()^3))))
+    m = @inferred mintegrate_exp(f1, mintegrate_exp(f2, mintegrate_exp(f3, mintegrate_exp(f4, StdUniform()^3))))
 
     for x in [Float32[0.7, 0.2, 0.5], Float32[-0.7, 0.2, 0.5]]
         @test @inferred(logdensityof(m, x)) isa Float32
@@ -229,13 +227,13 @@ end
 @testset "Density measures and Radon-Nikodym" begin
     x = randn()
     f(x) = x^2
-    @test log(𝒹(∫exp(f, Lebesgue()), Lebesgue())(x)) ≈ f(x)
+    @test log(density_rel(mintegrate_exp(f, Lebesgue()), Lebesgue())(x)) ≈ f(x)
 
-    let f = 𝒹(∫exp(x -> x^2, Lebesgue()), Lebesgue())
+    let f = density_rel(mintegrate_exp(x -> x^2, Lebesgue()), Lebesgue())
         @test log(f(x)) ≈ x^2
     end
 
-    let f = log𝒹(∫exp(x -> x^2, NormalMeasure()), NormalMeasure())
+    let f = logdensity_rel(mintegrate_exp(x -> x^2, NormalMeasure()), NormalMeasure())
         @test f(x) ≈ x^2
     end
 end

From f41a3bbdbaff0eb28b9ced27ad885768b6cf7b27 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 11:05:49 +0200
Subject: [PATCH 034/122] Add measure operators in submodule MeasureOperators
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The operators ⋄ (pushfwd), ⊙ (pullbck), ▷ (mbind), ⊗ (productmeasure),
∫ (mintegrate), ∫exp (mintegrate_exp), 𝒹 (density_rel) and log𝒹
(logdensity_rel) now live in the submodule `MeasureBase.MeasureOperators`,
so that users can opt into the operator syntax explicitly.

Ports the original measure-algebra commit 5404ff1.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/MeasureBase.jl        |   2 +
 src/measure_operators.jl  | 135 ++++++++++++++++++++++++++++++++++++++
 test/measure_operators.jl |  24 +++++++
 test/runtests.jl          |   2 +
 4 files changed, 163 insertions(+)
 create mode 100644 src/measure_operators.jl
 create mode 100644 test/measure_operators.jl

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 405114aa..8841488b 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -209,6 +209,8 @@ include("rand.jl")
 include("density.jl")
 include("density-core.jl")
 
+include("measure_operators.jl")
+
 include("interface.jl")
 
 using .Interface
diff --git a/src/measure_operators.jl b/src/measure_operators.jl
new file mode 100644
index 00000000..41606367
--- /dev/null
+++ b/src/measure_operators.jl
@@ -0,0 +1,135 @@
+"""
+    module MeasureOperators
+
+Defines the following operators for measures:
+
+* `f ⋄ μ == pushfwd(f, μ)`
+* `μ ⊙ f == pullbck(f, μ)`
+* `μ ▷ k == mbind(k, μ)`
+* `⊗(μs...) == productmeasure(μs)`
+* `∫(f, μ) == mintegrate(f, μ)`
+* `∫exp(f, μ) == mintegrate_exp(f, μ)`
+* `𝒹(ν, μ) == density_rel(ν, μ)`
+* `log𝒹(ν, μ) == logdensity_rel(ν, μ)`
+"""
+module MeasureOperators
+
+using MeasureBase: AbstractMeasure
+using MeasureBase: pushfwd, pullbck, mbind, productmeasure
+using MeasureBase: mintegrate, mintegrate_exp, density_rel, logdensity_rel
+using InverseFunctions: inverse
+
+@doc raw"""
+    ⋄(f, μ::AbstractMeasure) = pushfwd(f, μ)
+
+The `\\diamond` operator denotes a pushforward operation: `ν = f ⋄ μ`
+generates a
+[pushforward measure](https://en.wikipedia.org/wiki/Pushforward_measure).
+
+A common mathematical notation for a pushforward is ``f_*μ``, but as
+there is no "subscript-star" operator in Julia, we use `⋄`.
+
+See [`pushfwd(f, μ)`](@ref) for details.
+
+Also see [`ν ⊙ f`](@ref), the pullback operator.
+"""
+⋄(f, μ::AbstractMeasure) = pushfwd(f, μ)
+export ⋄
+
+@doc raw"""
+    ⊙(ν::AbstractMeasure, f) = pullbck(f, ν)
+
+The `\\odot` operator denotes a pullback operation.
+
+See also [`pullbck(ν, f)`](@ref) for details. Note that `pullbck` takes it's
+arguments in different order, in keeping with the Julia convention of
+passing functions as the first argument. A pullback is mathematically the
+precomposition of a measure `μ`` with the function `f` applied to sets. so
+`⊙` takes the measure as the first and the function as the second argument,
+as common in mathematical notation for precomposition.
+
+A common mathematical notation for pullback in measure theory is
+``f \circ μ``, but as `∘` is used for function composition in Julia and as
+`f` semantically acts point-wise on sets, we use `⊙`.
+
+Also see [f ⋄ μ](@ref), the pushforward operator.
+"""
+⊙(ν::AbstractMeasure, f) = pullbck(f, ν)
+export ⊙
+
+"""
+    μ ▷ k = mbind(k, μ)
+
+The `\\triangleright` operator denotes a measure monadic bind operation.
+
+A common operator choice for a monadic bind operator is `>>=` (e.g. in
+the Haskell programming language), but this has a different meaning in
+Julia and there is no close equivalent, so we use `▷`.
+
+See [`mbind(k, μ)`](@ref) for details. Note that `mbind` takes its
+arguments in different order, in keeping with the Julia convention of
+passing functions as the first argument. `▷`, on the other hand, takes
+its arguments in the order common for monadic binds in functional
+programming (like the Haskell `>>=` operator) and mathematics.
+"""
+▷(μ::AbstractMeasure, k) = mbind(k, μ)
+export ▷
+
+# ToDo: Use `⨂` instead of `⊗` for better readability?
+"""
+    ⊗(μs::AbstractMeasure...) = productmeasure(μs)
+
+`⊗` is an operator for building product measures.
+
+See [`productmeasure(μs)`](@ref) for details.
+"""
+⊗(μs::AbstractMeasure...) = productmeasure(μs)
+export ⊗
+
+"""
+    ∫(f, μ::AbstractMeasure) = mintegrate(f, μ)
+
+Denotes an indefinite integral of the function `f` with respect to the
+measure `μ`.
+
+See [`mintegrate(f, μ)`](@ref) for details.
+"""
+∫(f, μ::AbstractMeasure) = mintegrate(f, μ)
+export ∫
+
+"""
+    ∫exp(f, μ::AbstractMeasure) = mintegrate_exp(f, μ)
+
+Generates a new measure that is the indefinite integral of `exp` of `f`
+with respect to the measure `μ`.
+
+See [`mintegrate_exp(f, μ)`](@ref) for details.
+"""
+∫exp(f, μ::AbstractMeasure) = mintegrate_exp(f, μ)
+export ∫exp
+
+"""
+    𝒹(ν, μ) = density_rel(ν, μ)
+
+Compute the density, i.e. the
+[Radom-Nikodym derivative](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem)
+of `ν`` with respect to `μ`.
+
+For details, see [`density_rel(ν, μ)`}(@ref).
+"""
+𝒹(ν, μ::AbstractMeasure) = density_rel(ν, μ)
+export 𝒹
+
+"""
+    log𝒹(ν, μ) = logdensity_rel(ν, μ)
+
+Compute the log-density, i.e. the logarithm of the
+[Radom-Nikodym derivative](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem)
+of `ν`` with respect to `μ`.
+
+For details, see [`logdensity_rel(ν, μ)`}(@ref).
+"""
+log𝒹(ν, μ::AbstractMeasure) = logdensity_rel(ν, μ)
+export log𝒹
+
+end # module MeasureOperators
diff --git a/test/measure_operators.jl b/test/measure_operators.jl
new file mode 100644
index 00000000..1530f191
--- /dev/null
+++ b/test/measure_operators.jl
@@ -0,0 +1,24 @@
+using Test
+
+using MeasureBase: AbstractMeasure
+using MeasureBase: StdExponential, StdLogistic, StdNormal, StdUniform
+using MeasureBase: pushfwd, pullbck, mbind, productmeasure
+using MeasureBase: mintegrate, mintegrate_exp, density_rel, logdensity_rel
+using MeasureBase.MeasureOperators: ⋄, ⊙, ▷, ⊗, ∫, ∫exp, 𝒹, log𝒹
+
+@testset "MeasureOperators" begin
+    μ = StdExponential()
+    ν = StdUniform()
+    k(σ) = pushfwd(x -> σ * x, StdNormal())
+    μs = (StdExponential(), StdLogistic(), StdUniform())
+    f = sqrt
+
+    @test @inferred(f ⋄ μ) == pushfwd(f, μ)
+    @test @inferred(ν ⊙ f) == pullbck(f, ν)
+    @test @inferred(μ ▷ k) == mbind(k, μ)
+    @test @inferred(⊗(μs...)) == productmeasure(μs)
+    @test @inferred(∫(f, μ)) == mintegrate(f, μ)
+    @test @inferred(∫exp(f, μ)) == mintegrate_exp(f, μ)
+    @test @inferred(𝒹(ν, μ)) == density_rel(ν, μ)
+    @test @inferred(log𝒹(ν, μ)) == logdensity_rel(ν, μ)
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index f0d4488a..7183c58f 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -20,6 +20,8 @@ include("smf.jl")
 
 include("test_mooncake.jl")
 
+include("measure_operators.jl")
+
 include("combinators/weighted.jl")
 include("combinators/transformedmeasure.jl")
 include("combinators/reshape.jl")

From c5de45408c130a6ac8fd4fdbf8a792a6078026f1 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 11:08:55 +0200
Subject: [PATCH 035/122] Rework likelihoods around AbstractLikelihood
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Introduce the `AbstractLikelihood <: Function` supertype with the
`likelihood_kernel`/`likelihood_obs` accessor API and the
`as_likelihood`/`as_integrand` conversion hooks (including default
implementations for `Base.Fix2(densityof, x) ∘ f`-style objects, so
likelihood-like functions and densities compose with `mintegrate`).
`Likelihood` becomes a plain kernel/observation container; density
evaluation goes through the generic `AbstractLikelihood` methods and
converts kernel results via `asmeasure`.

Combines the original measure-algebra commits bf35d40, aa7416b, 470db23
and the likelihood-related parts of 1ed3f95.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/MeasureBase.jl            |   6 +-
 src/combinators/likelihood.jl | 289 +++++++++++++++++-----------------
 2 files changed, 151 insertions(+), 144 deletions(-)

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 8841488b..e9fcb312 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -176,6 +176,9 @@ include("primitive.jl")
 include("utils.jl")
 include("mass-interface.jl")
 
+include("density.jl")
+include("density-core.jl")
+
 include("primitives/counting.jl")
 include("primitives/lebesgue.jl")
 include("primitives/dirac.jl")
@@ -206,9 +209,6 @@ include("combinators/half.jl")
 
 include("rand.jl")
 
-include("density.jl")
-include("density-core.jl")
-
 include("measure_operators.jl")
 
 include("interface.jl")
diff --git a/src/combinators/likelihood.jl b/src/combinators/likelihood.jl
index 6dfd164f..40001007 100644
--- a/src/combinators/likelihood.jl
+++ b/src/combinators/likelihood.jl
@@ -1,207 +1,214 @@
-export AbstractLikelihood, Likelihood
+"""
+    abstract type AbstractLikelihood <: Function
 
-abstract type AbstractLikelihood end
+Abstract supertype for likelihood objects.
 
-# @inline function logdensityof(ℓ::AbstractLikelihood, p)
-#     t() = dynamic(unsafe_logdensityof(ℓ, p))
-#     f() = -Inf
-#     ifelse(insupport(ℓ, p), t, f)()
-# end
+Likelihoods are *not* measures, but density functions. They are callable
+and also support the DensityInterface API. If `ℒ isa AbstractLikelihood`,
+then
 
-# insupport(ℓ::AbstractLikelihood, p) = insupport(ℓ.k(p), ℓ.x)
+```julia
+DensityInterface.DensityKind(ℒ) == IsDensity()
+log(ℒ(θ)) ≈ logdensityof(ℒ, θ)
+```
 
-@doc raw"""
-    Likelihood(k::AbstractTransitionKernel, x)
+Given a transition kernel `k(θ)` (a function that takes a parameter object
+and returns a measure) and an observation `x`, the recommended way to create
+a likelihood object is
 
-"Observe" a value `x`, yielding a function from the parameters to ℝ.
+```julia
+ℒ = likelihoodof(k, x)
+ℒ isa AbstractLikelihood
+```
 
-Likelihoods are most commonly used in conjunction with an existing _prior_
-measure to yield a new measure, the _posterior_. In Bayes's Law, we have
+Then
 
-``P(θ|x) ∝ P(θ) P(x|θ)``
+```julia
+log(ℒ(θ)) ≈ logdensityof(ℒ, θ) ≈ logdensityof(k(θ), x)
+```
 
-Here ``P(θ)`` is the prior. If we consider ``P(x|θ)`` as a function on ``θ``,
-then it is called a likelihood.
+See [`likelihoodof`](@ref) for details on the mathematical semantics of
+`k` and `x`.
 
-Since measures are most commonly manipulated using `density` and `logdensity`,
-it's awkward to commit a (log-)likelihood to using one or the other. To evaluate
-a `Likelihood`, we therefore use `density` or `logdensity`, depending on the
-circumstances. In the latter case, it is of course acting as a log-density.
+Likelihood-like types that are not subtypes of `AbstractLikelihood` can be
+made compatible with the `MeasureBase` likelihoods and Lebesgue integrals by
+specializing [`MeasureBase.as_likelihood`](@ref) and
+[`MeasureBase.as_integrand`](@ref).
+"""
+abstract type AbstractLikelihood <: Function end
+export AbstractLikelihood
 
-For example,
+@inline AbstractLikelihood(l) = as_likelihood(l)::AbstractLikelihood
 
-    julia> ℓ = Likelihood(Normal{(:μ,)}, 2.0)
-    Likelihood(Normal{(:μ,), T} where T, 2.0)
+Base.convert(::Type{AbstractLikelihood}, l::AbstractLikelihood) = l
+Base.convert(::Type{AbstractLikelihood}, l) = AbstractLikelihood(l)
 
-    julia> density_def(ℓ, (μ=2.0,))
-    1.0
 
-    julia> logdensity_def(ℓ, (μ=2.0,))
-    -0.0
+"""
+    likelihood_kernel(ℒ::AbstractLikelihood)
 
-If, as above, the measure includes the parameter information, we can optionally
-leave it out of the second argument in the call to `density` or `logdensity`. 
+Return the transition kernel that is part of likelihood `ℒ`.
 
-    julia> density_def(ℓ, 2.0)
-    1.0
+If `ℒ = likelihoodof(k, x)` then `likelihood_kernel(ℒ)` must return an
+equivalent of `k` (typically but not necessarily `k` itself).
+"""
+function likelihood_kernel end
+export likelihood_kernel
 
-    julia> logdensity_def(ℓ, 2.0)
-    -0.0
 
-With several parameters, things work as expected:
-    
-    julia> ℓ = Likelihood(Normal{(:μ,:σ)}, 2.0)
-    Likelihood(Normal{(:μ, :σ), T} where T, 2.0)
-    
-    julia> logdensity_def(ℓ, (μ=2, σ=3))
-    -1.0986122886681098
-    
-    julia> logdensity_def(ℓ, (2,3))
-    -1.0986122886681098
-    
-    julia> logdensity_def(ℓ, [2, 3])
-    -1.0986122886681098
+"""
+    likelihood_obs(ℒ::AbstractLikelihood)
 
----------
+Return the observation that is part of likelihood `ℒ`.
 
-    Likelihood(M<:ParameterizedMeasure, constraint::NamedTuple, x)
+If `ℒ = likelihoodof(k, x)` then `likelihood_obs(ℒ)` must return an
+equivalent of `x` (typically but not necessarily `x` itself).
+"""
+function likelihood_obs end
+export likelihood_obs
 
-In some cases the measure might have several parameters, and we may want the
-(log-)likelihood with respect to some subset of them. In this case, we can use
-the three-argument form, where the second argument is a constraint. For example,
 
-    julia> ℓ = Likelihood(Normal{(:μ,:σ)}, (σ=3.0,), 2.0)
-    Likelihood(Normal{(:μ, :σ), T} where T, (σ = 3.0,), 2.0)
+"""
+    MeasureBase.as_likelihood(l)::AbstractLikelihood
 
-Similarly to the above, we have
+Turn a likelihood-like object `l` into an `AbstractLikelihood`.
 
-    julia> density_def(ℓ, (μ=2.0,))
-    0.3333333333333333
+Likelihood-like types that are not subtypes of `AbstractLikelihood` can be
+made compatible by specializing
 
-    julia> logdensity_def(ℓ, (μ=2.0,))
-    -1.0986122886681098
+```julia
+MeasureBase.as_likelihood(l::MyLikelihoodType) = likelihoodof(..., ...)
+MeasureBase.as_integrand(l::MyLikelihoodType) = MeasureBase.as_likelihood(l)
+```
 
-    julia> density_def(ℓ, 2.0)
-    0.3333333333333333
+By default, this is implemented for objects like
 
-    julia> logdensity_def(ℓ, 2.0)
-    -1.0986122886681098
+```julia
+l = Base.Fix2(densityof, x) ∘ f
+l = FuncDensity(Base.Fix2(densityof, x) ∘ f)
+l = LogFuncDensity(Base.Fix2(logdensityof, x) ∘ f)
+```
+"""
+function as_likelihood end
+export as_likelihood
 
------------------------
+@inline as_likelihood(l::AbstractLikelihood) = l
 
-Finally, let's return to the expression for Bayes's Law, 
+@inline as_integrand(l::AbstractLikelihood) = l
 
-``P(θ|x) ∝ P(θ) P(x|θ)``
 
-The product on the right side is computed pointwise. To work with this in
-MeasureBase, we have a "pointwise product" `⊙`, which takes a measure and a
-likelihood, and returns a new measure, that is, the unnormalized posterior that
-has density ``P(θ) P(x|θ)`` with respect to the base measure of the prior.
+(ℒ::AbstractLikelihood)(p) = densityof(ℒ, p)
 
-For example, say we have
 
-    μ ~ Normal()
-    x ~ Normal(μ,σ)
-    σ = 1
+DensityInterface.DensityKind(::AbstractLikelihood) = IsDensity()
 
-and we observe `x=3`. We can compute the posterior measure on `μ` as
 
-    julia> post = Normal() ⊙ Likelihood(Normal{(:μ, :σ)}, (σ=1,), 3)
-    Normal() ⊙ Likelihood(Normal{(:μ, :σ), T} where T, (σ = 1,), 3)
+_eval_k(ℒ::AbstractLikelihood, p) = asmeasure(likelihood_kernel(ℒ)(p))
 
-    julia> logdensity_def(post, 2)
-    -2.5
-"""
-struct Likelihood{K,X} <: AbstractLikelihood
-    k::K
-    x::X
+function DensityInterface.logdensityof(ℒ::AbstractLikelihood, p)
+    logdensityof(_eval_k(ℒ, p), likelihood_obs(ℒ))
+end
 
-    Likelihood(k::K, x::X) where {K<:AbstractTransitionKernel,X} = new{K,X}(k, x)
-    Likelihood(k::K, x::X) where {K<:Function,X} = new{K,X}(k, x)
-    Likelihood(μ, x) = Likelihood(kernel(μ), x)
+function DensityInterface.densityof(ℒ::AbstractLikelihood, p)
+    exp(ULogarithmic, logdensityof(_eval_k(ℒ, p), likelihood_obs(ℒ)))
 end
 
-(lik::AbstractLikelihood)(p) = exp(ULogarithmic, logdensityof(lik.k(p), lik.x))
 
-DensityInterface.DensityKind(::AbstractLikelihood) = IsDensity()
+const _SimpleLikelihood1 = ComposedFunction{<:Base.Fix2{typeof(densityof),<:Any},<:Any}
+as_likelihood(l::_SimpleLikelihood1) = likelihoodof(l.inner, l.outer.x)
+as_integrand(l::_SimpleLikelihood1) = as_likelihood(l)
 
-function Pretty.quoteof(ℓ::Likelihood)
-    k = Pretty.quoteof(ℓ.k)
-    x = Pretty.quoteof(ℓ.x)
-    :(Likelihood($k, $x))
-end
+const _SimpleLikelihood2 = DensityInterface.FuncDensity{
+    <:ComposedFunction{<:Base.Fix2{typeof(densityof),<:Any},<:Any},
+}
+as_likelihood(l::_SimpleLikelihood2) = likelihoodof(l._f.inner, l._f.outer.x)
+as_integrand(l::_SimpleLikelihood2) = as_likelihood(l)
 
-function Base.show(io::IO, ℓ::Likelihood)
-    io = IOContext(io, :compact => true)
-    Pretty.pprint(io, ℓ)
-end
+const _SimpleLikelihood3 = DensityInterface.LogFuncDensity{
+    <:ComposedFunction{<:Base.Fix2{typeof(logdensityof),<:Any},<:Any},
+}
+as_likelihood(l::_SimpleLikelihood3) = likelihoodof(l._log_f.inner, l._log_f.outer.x)
+as_integrand(l::_SimpleLikelihood3) = as_likelihood(l)
 
-insupport(ℓ::AbstractLikelihood, p) = insupport(ℓ.k(p), ℓ.x)
+const _SimpleLogLikelihood1 = ComposedFunction{<:Base.Fix2{typeof(logdensityof),<:Any},<:Any}
+as_integrand_exp(l::_SimpleLogLikelihood1) = likelihoodof(l.inner, l.outer.x)
 
-@inline function logdensityof(ℓ::AbstractLikelihood, p)
-    logdensityof(ℓ.k(p), ℓ.x)
-end
 
-@inline function unsafe_logdensityof(ℓ::AbstractLikelihood, p)
-    return unsafe_logdensityof(ℓ.k(p), ℓ.x)
-end
-
-# basemeasure(ℓ::Likelihood) = @error "Likelihood requires local base measure"
+@doc raw"""
+    struct Likelihood <: AbstractLikelihood
 
-export likelihoodof
+Default result of [`likelihoodof(k, x)`](@ref).
 
+See [`AbstractLikelihood`](@ref) and [`likelihoodof`](@ref) for details.
 """
-    likelihoodof(k::AbstractTransitionKernel, x; constraints...)
-    likelihoodof(k::AbstractTransitionKernel, x, constraints::NamedTuple)
+struct Likelihood{K,X} <: AbstractLikelihood
+    k::K
+    x::X
 
-A likelihood is *not* a measure. Rather, a likelihood acts on a measure, through
-the "pointwise product" `⊙`, yielding another measure.
-"""
-function likelihoodof end
+    Likelihood{K,X}(k, x) where {K,X} = new{K,X}(k, x)
+end
+export Likelihood
 
-likelihoodof(k, x, ::NamedTuple{()}) = Likelihood(k, x)
+# For type stability, in case k is a type (resp. a constructor):
+Likelihood(k, x::X) where {X} = Likelihood{Core.Typeof(k),X}(k, x)
 
-likelihoodof(k, x; kwargs...) = likelihoodof(k, x, NamedTuple(kwargs))
+likelihood_kernel(ℒ::Likelihood) = ℒ.k
+likelihood_obs(ℒ::Likelihood) = ℒ.x
 
-likelihoodof(k, x, pars::NamedTuple) = likelihoodof(kernel(k, pars), x)
+function Pretty.quoteof(ℒ::Likelihood)
+    k = Pretty.quoteof(ℒ.k)
+    x = Pretty.quoteof(ℒ.x)
+    :(Likelihood($k, $x))
+end
 
-likelihoodof(k::AbstractTransitionKernel, x) = Likelihood(k, x)
+function Base.show(io::IO, ℒ::Likelihood)
+    io = IOContext(io, :compact => true)
+    Pretty.pprint(io, ℒ)
+end
 
-export log_likelihood_ratio
 
-"""
-    log_likelihood_ratio(ℓ::Likelihood, p, q)
+@doc raw"""
+    likelihoodof(k, x)::AbstractLikelihood
 
-Compute the log of the likelihood ratio, in order to compare two choices for
-parameters. This is computed as
+Returns the likelihood of observing `x` under a family of probability
+measures that is generated by a transition kernel `k(θ)`.
 
-    logdensity_rel(ℓ.k(p), ℓ.k(q), ℓ.x)
+`k(θ)` maps points in the parameter space to measures (resp. objects that can
+be converted to measures) on an implicit set `Χ` that contains values like
+`x`.
 
-Since `logdensity_rel` can leave common base measure unevaluated, this can be
-more efficient than
+`likelihoodof(k, x)` returns a likelihood object. A likelihood is **not** a
+measure, it is a function from the parameter space to `ℝ₊`. Likelihood
+objects can also be interpreted as "generic densities" (but **not** as
+probability densities).
 
-    logdensityof(ℓ.k(p), ℓ.x) - logdensityof(ℓ.k(q), ℓ.x)
-"""
-log_likelihood_ratio(ℓ::Likelihood, p, q) = logdensity_rel(ℓ.k(p), ℓ.k(q), ℓ.x)
+`likelihoodof(k, x)` implicitly chooses `ξ = rootmeasure(k(θ))` as the
+reference measure on the observation set `Χ`. Note that this implicit
+`ξ` **must** be independent of `θ`.
 
-# likelihoodof(k, x; kwargs...) = likelihoodof(k, x, NamedTuple(kwargs))
+`ℒ = likelihoodof(k, x)` has the mathematical interpretation
 
-export likelihood_ratio
+```math
+\mathcal{L}_x(\theta) = \frac{\rm{d}\, k(\theta)}{\rm{d}\, \chi}(x)
+```
 
-"""
-    likelihood_ratio(ℓ::Likelihood, p, q)
+`likelihoodof` must return an object that implements the
+[`DensityInterface`](https://github.com/JuliaMath/DensityInterface.jl) API
+and `ℒ = likelihoodof(k, x)` must satisfy
 
-Compute the log of the likelihood ratio, in order to compare two choices for
-parameters. This is equal to
+```julia
+log(ℒ(θ)) == logdensityof(ℒ, θ) ≈ logdensityof(k(θ), x)
 
-    density_rel(ℓ.k(p), ℓ.k(q), ℓ.x)
+DensityKind(ℒ) isa IsDensity
+```
 
-but is computed using LogarithmicNumbers.jl to avoid underflow and overflow.
-Since `density_rel` can leave common base measure unevaluated, this can be
-more efficient than
+[`likelihood_kernel(ℒ)`](@ref) must return an equivalent of `k` and
+[`likelihood_obs(ℒ)`](@ref) must return an equivalent of `x` (typically, but
+not necessarily, `k` and `x` themselves).
 
-    logdensityof(ℓ.k(p), ℓ.x) - logdensityof(ℓ.k(q), ℓ.x)
+By default, an instance of [`MeasureBase.Likelihood`](@ref) is returned.
 """
-function likelihood_ratio(ℓ::Likelihood, p, q)
-    exp(ULogarithmic, logdensity_rel(ℓ.k(p), ℓ.k(q), ℓ.x))
-end
+function likelihoodof end
+export likelihoodof
+
+likelihoodof(k, x) = Likelihood(k, x)

From 2fb782e9ce4d3ce0781e30776f3892377d6d07e8 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 11:09:29 +0200
Subject: [PATCH 036/122] Remove remnants of pointwiseproduct

Ports the original measure-algebra commit 28c6d30.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/combinators/smart-constructors.jl | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index 26ba3948..45b1c5fa 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -4,18 +4,6 @@
 
 half(μ::AbstractMeasure) = Half(μ)
 
-###############################################################################
-# PointwiseProductMeasure
-
-function pointwiseproduct(μ::AbstractMeasure, ℓ::Likelihood)
-    T = Core.Compiler.return_type(ℓ.k, Tuple{gentype(μ)})
-    return pointwiseproduct(T, μ, ℓ)
-end
-
-function pointwiseproduct(::Type{T}, μ::AbstractMeasure, ℓ::Likelihood) where {T}
-    return PointwiseProductMeasure(μ, ℓ)
-end
-
 ###############################################################################
 # PowerMeaure
 

From 15a0b9ec0f0b53fac14c32839cccb94fa14f2842 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 11:09:52 +0200
Subject: [PATCH 037/122] Remove splat (unused here, and Base has it now)

Ports the original measure-algebra commit 81a7d93.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/MeasureBase.jl |  1 -
 src/splat.jl       | 11 -----------
 2 files changed, 12 deletions(-)
 delete mode 100644 src/splat.jl

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index e9fcb312..e7c9f8a8 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -167,7 +167,6 @@ include("smf.jl")
 include("getdof.jl")
 include("transport.jl")
 include("schema.jl")
-include("splat.jl")
 include("proxies.jl")
 include("kernel.jl")
 include("parameterized.jl")
diff --git a/src/splat.jl b/src/splat.jl
deleted file mode 100644
index d1df4f17..00000000
--- a/src/splat.jl
+++ /dev/null
@@ -1,11 +0,0 @@
-struct Splat{F}
-    f::F
-end
-
-function (s::Splat{F})(x) where {F}
-    s.f(x...)
-end
-
-unsplat(s::Splat) = s.f
-
-splat(f) = Splat(f)

From 5bc1389d1fcf4fdc4fd5a166051912dcaf4e2671 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 12:53:54 +0200
Subject: [PATCH 038/122] Add fast_dof, some_dof and NoFastInsupport

Introduce the `AbstractNoDOF` type hierarchy with absorbing arithmetic,
`NoFastDOF` and `fast_dof` (defaults to `getdof`, to be specialized by
measures whose DOF can't be computed instantly) plus `some_dof` (DOF at
an unspecified point, via `localmeasure`). `_default_getdof` now returns
a `NoDOF` instance instead of the type, and `check_dof` uses `fast_dof`
and tolerates unknown DOF. Add `NoFastInsupport` and make
`require_insupport` tolerate it, so measures like monadic binds can
declare that support checking is not cheap.

Ports parts of the original measure-algebra commits 7e93050 and later
STASH fixes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/getdof.jl    | 95 +++++++++++++++++++++++++++++++++++++++++++++---
 src/insupport.jl | 27 +++++++++++---
 2 files changed, 110 insertions(+), 12 deletions(-)

diff --git a/src/getdof.jl b/src/getdof.jl
index 16ae7cc6..c2edd15c 100644
--- a/src/getdof.jl
+++ b/src/getdof.jl
@@ -1,11 +1,30 @@
 """
-    MeasureBase.NoDOF{MU}
+    abstract type MeasureBase.AbstractNoDOF{MU}
+
+Abstract supertype for [`NoDOF`](@ref) and [`NoFastDOF`](@ref).
+"""
+abstract type AbstractNoDOF{MU} end
+
+Base.:+(nodof::AbstractNoDOF) = nodof
+Base.:+(::IntegerLike, nodof::AbstractNoDOF) = nodof
+Base.:+(nodof::AbstractNoDOF, ::IntegerLike) = nodof
+Base.:+(nodof::AbstractNoDOF, ::AbstractNoDOF) = nodof
+
+Base.:*(nodof::AbstractNoDOF) = nodof
+Base.:*(::IntegerLike, nodof::AbstractNoDOF) = nodof
+Base.:*(nodof::AbstractNoDOF, ::IntegerLike) = nodof
+Base.:*(nodof::AbstractNoDOF, ::AbstractNoDOF) = nodof
+
+
+"""
+    MeasureBase.NoDOF{MU} <: AbstractNoDOF{MU}
 
 Indicates that there is no way to compute degrees of freedom of a measure
 of type `MU` with the given information, e.g. because the DOF are not
 a global property of the measure.
 """
-struct NoDOF{MU} end
+struct NoDOF{MU} <: AbstractNoDOF{MU} end
+
 
 """
     getdof(μ)
@@ -22,22 +41,86 @@ Also see [`check_dof`](@ref).
 function getdof end
 
 # Prevent infinite recursion:
-@inline _default_getdof(::Type{MU}, ::MU) where {MU} = NoDOF{MU}
+@inline _default_getdof(::Type{MU}, ::MU) where {MU} = NoDOF{MU}()
 @inline _default_getdof(::Type{MU}, mu_base) where {MU} = getdof(mu_base)
 
 @inline getdof(μ::MU) where {MU} = _default_getdof(MU, basemeasure(μ))
 
+
+"""
+    MeasureBase.NoFastDOF{MU} <: AbstractNoDOF{MU}
+
+Indicates that there is no way to compute the degrees of freedom of a
+measure of type `MU` efficiently.
+"""
+struct NoFastDOF{MU} <: AbstractNoDOF{MU} end
+
+
+"""
+    fast_dof(μ::MU)
+
+Returns the effective number of degrees of freedom of variates of
+measure `μ`, if it can be computed efficiently, otherwise
+returns [`NoFastDOF{MU}()`](@ref).
+
+Defaults to `getdof(μ)` and should be specialized for measures for
+which DOF can't be computed instantly.
+
+Also see [`getdof`](@ref) and [`check_dof`](@ref).
+"""
+function fast_dof end
+export fast_dof
+
+fast_dof(μ) = getdof(μ)
+
+
+"""
+    MeasureBase.some_dof(μ::AbstractMeasure)
+
+Get the DOF at some unspecified point of measure `μ`.
+
+Use with caution!
+
+In general, use [`getdof(μ)`](@ref) instead. `some_dof` is useful for
+measures that are expected to have a constant DOF over their whole
+space, but for which there is no way to compute it (or prove that
+the DOF is constant over the measurable space).
+"""
+function some_dof end
+
+function some_dof(μ)
+    m = asmeasure(μ)
+    _try_direct_dof(m, getdof(m))
+end
+
+_try_direct_dof(::AbstractMeasure, dof::IntegerLike) = dof
+_try_direct_dof(μ::AbstractMeasure, ::AbstractNoDOF) =
+    _try_local_dof(μ, some_dof(_some_localmeasure(μ)))
+
+_try_local_dof(::AbstractMeasure, dof::IntegerLike) = dof
+_try_local_dof(μ::AbstractMeasure, ::AbstractNoDOF) =
+    throw(ArgumentError("Can't determine DOF for measure of type $(nameof(typeof(μ)))"))
+
+_some_localmeasure(μ::AbstractMeasure) = localmeasure(μ, testvalue(μ))
+
+
 """
     MeasureBase.check_dof(ν, μ)::Nothing
 
 Check if `ν` and `μ` have the same effective number of degrees of freedom
-according to [`MeasureBase.getdof`](@ref).
+according to [`MeasureBase.fast_dof`](@ref).
+
+Does not throw an exception if the DOF of `ν` or `μ` can't be computed
+efficiently.
 """
 function check_dof end
 
 function check_dof(ν, μ)
-    n_ν = getdof(ν)
-    n_μ = getdof(μ)
+    n_ν = fast_dof(ν)
+    n_μ = fast_dof(μ)
+    if n_ν isa AbstractNoDOF || n_μ isa AbstractNoDOF
+        return nothing
+    end
     if n_ν != n_μ
         throw(
             ArgumentError(
diff --git a/src/insupport.jl b/src/insupport.jl
index a9a96363..bb09b1e6 100644
--- a/src/insupport.jl
+++ b/src/insupport.jl
@@ -1,12 +1,21 @@
 """
-    inssupport(m, x)
+    MeasureBase.NoFastInsupport{MU}
+
+Indicates that there is no fast way to compute if a point lies within the
+support of measures of type `MU`.
+"""
+struct NoFastInsupport{MU} end
+
+
+"""
+    insupport(m, x)
     insupport(m)
 
-`insupport(m,x)` computes whether `x` is in the support of `m`.
+`insupport(m, x)` computes whether `x` is in the support of `m` and
+returns either a `Bool` or an instance of [`NoFastInsupport`](@ref).
 
 `insupport(m)` returns a function, and satisfies
-
-insupport(m)(x) == insupport(m, x)
+`insupport(m)(x) == insupport(m, x)`.
 """
 function insupport end
 
@@ -15,12 +24,18 @@ function insupport end
 
 Checks if `x` is in the support of distribution/measure `μ`, throws an
 `ArgumentError` if not.
+
+Will not throw an exception if `insupport` returns an instance of
+[`NoFastInsupport`](@ref).
 """
 function require_insupport end
 
 function require_insupport(μ, x)
-    if !insupport(μ, x)
-        throw(ArgumentError("x is not within the support of μ"))
+    ins = insupport(μ, x)
+    if !(ins isa NoFastInsupport)
+        if !ins
+            throw(ArgumentError("x is not within the support of μ"))
+        end
     end
     return nothing
 end

From c123ef05da26200a92582cfee02b68c73d0ff01e Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 13:00:04 +0200
Subject: [PATCH 039/122] Add localmeasure and transportmeasure

Introduce the `localmeasure`/`transportmeasure` interface: both default
to the measure itself and return a measure that behaves like the
original one in the infinitesimal neighborhood of a point, for density
calculation resp. density calculation and transport. Measures like
monadic binds specialize them. `unsafe_logdensityof` and
`unsafe_logdensity_rel` now evaluate via the local measure, with an
extra dispatch boundary to reduce the number of required
specializations. `logdensity_rel` handles `NoFastInsupport`, and
`_checksupport` passes results through unchanged in that case.

Transport now converts its arguments via `asmeasure`, uses `fast_dof`
for the standard intermediate measure, guards against overly deep
transport-origin stacks, and powers of `NoTransportOrigin` stay
`NoTransportOrigin`. The two-argument form of `basemeasure` is gone.

Ports parts of the original measure-algebra commits fe28297, 95632f4,
67cf8ea and related STASH commits.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/density-core.jl | 93 ++++++++++++++++++++++++++++++++++++++++++---
 src/interface.jl    |  2 +-
 src/transport.jl    | 23 +++++++++--
 src/utils.jl        |  2 -
 4 files changed, 109 insertions(+), 11 deletions(-)

diff --git a/src/density-core.jl b/src/density-core.jl
index f3b2db2b..33a3c3cb 100644
--- a/src/density-core.jl
+++ b/src/density-core.jl
@@ -38,6 +38,40 @@ end
 end
 
 _checksupport(cond, result) = ifelse(cond == true, result, oftype(result, -Inf))
+@inline _checksupport(::NoFastInsupport, result) = result
+
+
+"""
+    localmeasure(m::AbstractMeasure, x)::AbstractMeasure
+
+Return a measure that behaves like `m` in the infinitesimal neighborhood
+of `x` in respect to density calculation.
+
+Note that the resulting measure may not be well defined outside of the
+infinitesimal neighborhood of `x`.
+
+For most measure types simply returns `m` itself. [`mbind`](@ref),
+for example, generates measures for which `localmeasure(m, x)` depends
+on `x`.
+"""
+localmeasure(m::AbstractMeasure, x) = m
+export localmeasure
+
+
+"""
+    MeasureBase.transportmeasure(m::AbstractMeasure, x)::AbstractMeasure
+
+Return a measure that behaves like `m` in the infinitesimal neighborhood
+of `x` in respect to both transport and density calculation.
+
+Note that the resulting measure may not be well defined outside of the
+infinitesimal neighborhood of `x`.
+
+For most measure types simply returns `m` itself. [`mbind`](@ref),
+for example, generates measures for which `transportmeasure(m, x)` depends
+on `x`.
+"""
+transportmeasure(m::AbstractMeasure, x) = m
 
 export unsafe_logdensityof
 
@@ -50,11 +84,17 @@ This is "unsafe" because it does not check `insupport(m, x)`.
 
 See also `logdensityof`.
 """
-@inline function unsafe_logdensityof(μ::M, x) where {M}
+@inline function unsafe_logdensityof(μ::AbstractMeasure, x)
+    μ_local = localmeasure(μ, x)
+    # Extra dispatch boundary to reduce number of required specializations of implementation:
+    return _unsafe_logdensityof_local(μ_local, x)
+end
+
+@inline function _unsafe_logdensityof_local(μ::M, x) where {M}
     ℓ_0 = logdensity_def(μ, x)
     b_0 = μ
     Base.Cartesian.@nexprs 10 i -> begin  # 10 is just some "big enough" number
-        b_{i} = basemeasure(b_{i - 1}, x)
+        b_{i} = basemeasure(b_{i - 1})
 
         # The below makes the evaluated code shorter, but screws up Zygote
         # if b_{i} isa typeof(b_{i - 1})
@@ -76,20 +116,56 @@ known to be in the support of both, it can be more efficient to call
 `unsafe_logdensity_rel`. 
 """
 @inline function logdensity_rel(μ::M, ν::N, x::X) where {M,N,X}
+    inμ = insupport(μ, x)
+    inν = insupport(ν, x)
+    return _logdensity_rel_impl(μ, ν, x, inμ, inν)
+end
+
+@inline function _logdensity_rel_impl(μ::M, ν::N, x::X, inμ::Bool, inν::Bool) where {M,N,X}
     T = unstatic(
         promote_type(
             return_type(logdensity_def, (μ, x)),
             return_type(logdensity_def, (ν, x)),
         ),
     )
-    inμ = insupport(μ, x)
-    inν = insupport(ν, x)
     istrue(inμ) || return convert(T, ifelse(inν, -Inf, NaN))
     istrue(inν) || return convert(T, Inf)
 
     return unsafe_logdensity_rel(μ, ν, x)
 end
 
+@inline function _logdensity_rel_impl(
+    μ::M,
+    ν::N,
+    x::X,
+    @nospecialize(::NoFastInsupport),
+    @nospecialize(::NoFastInsupport)
+) where {M,N,X}
+    unsafe_logdensity_rel(μ, ν, x)
+end
+
+@inline function _logdensity_rel_impl(
+    μ::M,
+    ν::N,
+    x::X,
+    inμ::Bool,
+    @nospecialize(::NoFastInsupport)
+) where {M,N,X}
+    logd = unsafe_logdensity_rel(μ, ν, x)
+    return istrue(inμ) ? logd : oftype(logd, -Inf)
+end
+
+@inline function _logdensity_rel_impl(
+    μ::M,
+    ν::N,
+    x::X,
+    @nospecialize(::NoFastInsupport),
+    inν::Bool
+) where {M,N,X}
+    logd = unsafe_logdensity_rel(μ, ν, x)
+    return istrue(inν) ? logd : oftype(logd, +Inf)
+end
+
 """
     unsafe_logdensity_rel(m1, m2, x)
 
@@ -98,7 +174,14 @@ known to be in the support of both `m1` and `m2`.
 
 See also `logdensity_rel`.
 """
-@inline function unsafe_logdensity_rel(μ::M, ν::N, x::X) where {M,N,X}
+@inline function unsafe_logdensity_rel(μ::AbstractMeasure, ν::AbstractMeasure, x)
+    μ_local = localmeasure(μ, x)
+    ν_local = localmeasure(ν, x)
+    # Extra dispatch boundary to reduce number of required specializations of implementation:
+    return _unsafe_logdensity_rel_local(μ_local, ν_local, x)
+end
+
+@inline function _unsafe_logdensity_rel_local(μ::M, ν::N, x::X) where {M,N,X}
     if static_hasmethod(logdensity_def, Tuple{M,N,X})
         return logdensity_def(μ, ν, x)
     end
diff --git a/src/interface.jl b/src/interface.jl
index 4890ddd6..6003203d 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -64,7 +64,7 @@ function test_interface(μ::M) where {M}
             # testvalue, logdensityof
 
             x = @inferred testvalue(Float64, μ)
-            β = @inferred basemeasure(μ, x)
+            β = @inferred basemeasure(μ)
 
             ℓμ = @inferred logdensityof(μ, x)
             ℓβ = @inferred logdensityof(β, x)
diff --git a/src/transport.jl b/src/transport.jl
index b0c8ed41..9cea6fed 100644
--- a/src/transport.jl
+++ b/src/transport.jl
@@ -8,6 +8,8 @@ See [`MeasureBase.transport_origin`](@ref).
 """
 struct NoTransportOrigin{NU} end
 
+Base.:^(origin::NoTransportOrigin, ::IntegerLike) = origin
+
 """
     MeasureBase.transport_origin(ν)
 
@@ -147,6 +149,21 @@ end
     μ,
     x,
 ) where {n_ν,n_μ}
+    if n_ν == 10
+        return :(throw(
+            ArgumentError(
+                "Transport to measure of type $(nameof(typeof(ν))) not supported, origin stack too deep.",
+            ),
+        ))
+    end
+    if n_μ == 10
+        return :(throw(
+            ArgumentError(
+                "Transport from measure of type $(nameof(typeof(μ))) not supported, origin stack too deep.",
+            ),
+        ))
+    end
+
     prog = quote
         μ0 = μ
         x0 = x
@@ -183,8 +200,8 @@ end
     return prog
 end
 
-@inline _transport_intermediate(ν, μ) = _transport_intermediate(getdof(ν), getdof(μ))
-@inline _transport_intermediate(::Integer, n_μ::Integer) = StdUniform()^n_μ
+@inline _transport_intermediate(ν, μ) = _transport_intermediate(fast_dof(ν), fast_dof(μ))
+@inline _transport_intermediate(::IntegerLike, n_μ::IntegerLike) = StdUniform()^n_μ
 @inline _transport_intermediate(::StaticInteger{1}, ::StaticInteger{1}) = StdUniform()
 
 _call_transport_def(ν, μ, x) = transport_def(ν, μ, x)
@@ -227,7 +244,7 @@ struct TransportFunction{NU,MU} <: Function
     end
 end
 
-@inline transport_to(ν, μ) = TransportFunction(ν, μ)
+@inline transport_to(ν, μ) = TransportFunction(asmeasure(ν), asmeasure(μ))
 
 function Base.:(==)(a::TransportFunction, b::TransportFunction)
     return a.ν == b.ν && a.μ == b.μ
diff --git a/src/utils.jl b/src/utils.jl
index c1e97034..e169c7c1 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -19,8 +19,6 @@ testvalue(::Type{T}) where {T} = zero(T)
 
 export rootmeasure
 
-basemeasure(μ, x) = basemeasure(μ)
-
 """
     rootmeasure(μ::AbstractMeasure)
 

From 596593c8126e0ee7663f2e3ec0114a83dd6e298f Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 13:02:35 +0200
Subject: [PATCH 040/122] Add collection utils, StdPowerMeasure and power
 measure fast_dof

Add the collection helpers `_as_tuple`, `_get_or_view`, `_split_after`
(vector/tuple/NamedTuple), `_fill_value`/`_fill_axes` and
`_flatten_to_rv`, in preparation for the generalized product transport
and combined measures. Add the `StdPowerMeasure` type alias, `fast_dof`
for power measures and make power measure `insupport` propagate
`NoFastInsupport` from the parent measure.

Ports parts of the original measure-algebra commits 763bd70, cd2db56
and 7e93050.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/MeasureBase.jl         |  3 ++
 src/collection_utils.jl    | 72 ++++++++++++++++++++++++++++++++++++++
 src/combinators/power.jl   | 15 +++++++-
 src/standard/stdmeasure.jl |  7 ++++
 4 files changed, 96 insertions(+), 1 deletion(-)

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index e7c9f8a8..2aa92089 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -60,6 +60,9 @@ using StaticThings:
 import HeterogeneousComputing
 using HeterogeneousComputing: real_numtype
 
+using ArraysOfArrays:
+    VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays, VectorOfSimilarVectors, flatview
+
 export gentype
 
 export AbstractMeasure
diff --git a/src/collection_utils.jl b/src/collection_utils.jl
index 1de51f7e..1ef9f9fd 100644
--- a/src/collection_utils.jl
+++ b/src/collection_utils.jl
@@ -22,3 +22,75 @@ _rev_cumsum(xs::AbstractVector) = reverse(cumsum(reverse(xs)))
 
 # Equivalent to `cumprod(xs)``:
 _exp_cumsum_log(xs::AbstractVector) = exp.(cumsum(log.(xs)))
+
+Base.@propagate_inbounds _as_tuple(v::AbstractVector, ::Val{N}) where {N} = Tuple(SVector{N}(v))
+
+
+Base.@propagate_inbounds function _get_or_view(A::AbstractVector, from::IntegerLike, until::IntegerLike)
+    view(A, from:until)
+end
+
+Base.@propagate_inbounds function _get_or_view(
+    A::AbstractVector,
+    ::StaticInteger{from},
+    ::StaticInteger{until},
+) where {from,until}
+    SVector{until - from + 1}(view(A, from:until))
+end
+
+# ToDo: Specialize for StaticVector instead of SVector?
+Base.@propagate_inbounds function _get_or_view(
+    A::SVector,
+    from::StaticInteger,
+    until::StaticInteger,
+)
+    # ToDo: Improve implementation:
+    SVector(_get_or_view(Tuple(A), from, until))
+end
+
+Base.@propagate_inbounds function _get_or_view(tpl::Tuple, from::IntegerLike, until::IntegerLike)
+    ntuple(i -> tpl[from + i - 1], Val(until - from + 1))
+end
+
+
+@inline function _split_after(x::AbstractVector, n::IntegerLike)
+    idxs = maybestatic_eachindex(x)
+    i_first = maybestatic_first(idxs)
+    i_last = maybestatic_last(idxs)
+    _get_or_view(x, i_first, i_first + n - one(n)), _get_or_view(x, i_first + n, i_last)
+end
+
+@inline _split_after(x::Tuple, n) = _split_after(x::Tuple, Val{n}())
+@inline _split_after(x::Tuple, ::Val{N}) where {N} = x[begin:(begin+N-1)], x[(begin+N):end]
+
+@generated function _split_after(x::NamedTuple{names}, ::Val{names_a}) where {names,names_a}
+    n = length(names_a)
+    if names[begin:(begin+n-1)] == names_a
+        names_b = names[(begin+n):end]
+        quote
+            a, b = _split_after(values(x), Val($n))
+            NamedTuple{$names_a}(a), NamedTuple{$names_b}(b)
+        end
+    else
+        quote
+            throw(ArgumentError("Can't split NamedTuple{$names} after {$names_a}"))
+        end
+    end
+end
+
+
+# Field access functions for Fill:
+_fill_value(x::FillArrays.Fill) = x.value
+_fill_axes(x::FillArrays.Fill) = x.axes
+
+
+_flatten_to_rv(VV::AbstractVector{<:AbstractVector{<:Real}}) = flatview(VectorOfArrays(VV))
+_flatten_to_rv(VV::AbstractVector{<:StaticVector{N,<:Real}}) where {N} =
+    flatview(VectorOfSimilarArrays(VV))
+
+_flatten_to_rv(VV::VectorOfSimilarVectors{<:Real}) = flatview(VV)
+_flatten_to_rv(VV::VectorOfVectors{<:Real}) = flatview(VV)
+
+_flatten_to_rv(::Tuple{}) = []
+_flatten_to_rv(tpl::Tuple{Vararg{AbstractVector}}) = vcat(tpl...)
+_flatten_to_rv(tpl::Tuple{Vararg{StaticVector}}) = vcat(tpl...)
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 6f065c3f..7e22fb36 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -144,15 +144,28 @@ end
     end
 end
 
+_all(A) = all(A)
+_all(::AbstractArray{NoFastInsupport{T}}) where {T} = NoFastInsupport{T}()
+
 @inline function insupport(μ::PowerMeasure, x::AbstractArray)
     p = μ.parent
-    all(x) do xj
+    insupp = broadcast(x) do xj
         # https://github.com/SciML/Static.jl/issues/36
         dynamic(insupport(p, xj))
     end
+    _all(insupp)
 end
 
 @inline getdof(μ::PowerMeasure) = getdof(μ.parent) * size2length(axes2size(μ.axes))
+@inline fast_dof(μ::PowerMeasure) = fast_dof(μ.parent) * size2length(axes2size(μ.axes))
+
+# Static.SOneTo(0) is not static (yet):
+@inline function getdof(::PowerMeasure{<:Any,<:NTuple{N,StaticOneToLike{0}}}) where {N}
+    static(0)
+end
+@inline function fast_dof(::PowerMeasure{<:Any,<:NTuple{N,StaticOneToLike{0}}}) where {N}
+    static(0)
+end
 
 @propagate_inbounds function checked_arg(μ::PowerMeasure, x::AbstractArray{<:Any})
     @boundscheck begin
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index a9c09b5c..3dd30e24 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -4,6 +4,13 @@ StdMeasure(::typeof(rand)) = StdUniform()
 StdMeasure(::typeof(randexp)) = StdExponential()
 StdMeasure(::typeof(randn)) = StdNormal()
 
+"""
+    MeasureBase.StdPowerMeasure{MU<:StdMeasure,N}
+
+The type of an `N`-dimensional power of a standard measure of type `MU`.
+"""
+const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}}
+
 @inline check_dof(::StdMeasure, ::StdMeasure) = nothing
 
 @inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x

From c97bd3d872128a6dec690f04e55fe84f1f4871bb Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 13:07:57 +0200
Subject: [PATCH 041/122] Canonical measure nesting in smart constructors

Rework the `powermeasure` and `productmeasure` smart constructors to
maintain a canonical measure type nesting (WeightedMeasure over Dirac
over PowerMeasure over ProductMeasure): powers and products of Dirac
measures become Dirac measures of filled values, weights are pulled out
of powers, marginal collections are converted via `asmeasure`, products
over `Fill` and over singleton-typed arrays collapse to power measures.
Product measures also get `fast_dof` and `NoFastInsupport`-aware
`insupport`, plus a power-measure proxy for `Fill` marginals.
`powermeasure` now accepts sizes, axes and plain integer exponents.

Ports the original measure-algebra commits 00b4d61, b28a902 and parts
of related STASH commits.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/combinators/power.jl              |   4 -
 src/combinators/product.jl            |  11 ++-
 src/combinators/smart-constructors.jl | 106 +++++++++++++++++++++-----
 3 files changed, 96 insertions(+), 25 deletions(-)

diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 7e22fb36..1ce9ca98 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -75,10 +75,6 @@ function Base.rand(rng::AbstractRNG, ::Type{T}, d::PowerMeasure) where {T}
     end
 end
 
-@inline function powermeasure(x::T, sz::Tuple{Vararg{Any,N}}) where {T,N}
-    PowerMeasure(x, asaxes(sz))
-end
-
 marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes)
 
 function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{AbstractArray,N}}) where {N}
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index dbdfffa3..468c38c0 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -28,6 +28,9 @@ Base.size(μ::AbstractProductMeasure) = size(marginals(μ))
 
 basemeasure(d::AbstractProductMeasure) = productmeasure(map(basemeasure, marginals(d)))
 
+proxy(μ::ProductMeasure{<:FillArrays.Fill}) =
+    powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ)))
+
 function Base.rand(rng::AbstractRNG, ::Type{T}, d::AbstractProductMeasure) where {T}
     mar = marginals(d)
     _rand_product(rng, T, mar, eltype(mar))
@@ -221,12 +224,16 @@ end
 
 @inline function insupport(d::AbstractProductMeasure, x)
     for (mj, xj) in zip(marginals(d), x)
-        dynamic(insupport(mj, xj)) || return false
+        insup = dynamic(insupport(mj, xj))
+        if insup isa NoFastInsupport || insup == false
+            return insup
+        end
     end
     return true
 end
 
-getdof(d::AbstractProductMeasure) = mapreduce(getdof, +, marginals(d))
+getdof(d::AbstractProductMeasure) = sum(getdof, marginals(d))
+fast_dof(d::AbstractProductMeasure) = sum(fast_dof, marginals(d))
 
 function checked_arg(μ::ProductMeasure{<:NTuple{N,Any}}, x::NTuple{N,Any}) where {N}
     map(checked_arg, marginals(μ), x)
diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index 45b1c5fa..cb15f6ad 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -1,4 +1,9 @@
 
+# Canonical measure type nesting, outer to inner:
+#
+# WeightedMeasure, Dirac, PowerMeasure, ProductMeasure
+
+
 ###############################################################################
 # Half
 
@@ -7,44 +12,107 @@ half(μ::AbstractMeasure) = Half(μ)
 ###############################################################################
 # PowerMeaure
 
-powermeasure(m::AbstractMeasure, ::Tuple{}) = m
+"""
+    powermeasure(μ, dims)
+    powermeasure(μ, axes)
+
+Constructs a power of a measure `μ`.
+
+`powermeasure(μ, exponent)` is semantically equivalent to
+`productmeasure(Fill(μ, exponent))`, but more efficient.
+"""
+function powermeasure end
+export powermeasure
+
+@inline powermeasure(μ, exponent) = _generic_powermeasure_stage1(asmeasure(μ), asaxes(exponent))
+
+@inline _generic_powermeasure_stage1(μ::AbstractMeasure, ::Tuple{}) = μ
 
-function powermeasure(
-    μ::WeightedMeasure,
-    dims::Tuple{<:AbstractArray,Vararg{AbstractArray}},
-)
-    k = mapreduce(length, *, dims) * μ.logweight
-    return weightedmeasure(k, μ.base^dims)
+@inline function _generic_powermeasure_stage1(μ::AbstractMeasure, exponent::Tuple)
+    _generic_powermeasure_stage2(μ, exponent)
 end
 
-function powermeasure(μ::WeightedMeasure, dims::NonEmptyTuple)
-    k = prod(dims) * μ.logweight
-    return weightedmeasure(k, μ.base^dims)
+@inline _generic_powermeasure_stage2(μ::AbstractMeasure, exponent::Tuple) =
+    PowerMeasure(μ, exponent)
+
+@inline function _generic_powermeasure_stage2(μ::Dirac, exponent::Tuple)
+    Dirac(maybestatic_fill(μ.x, exponent))
+end
+
+@inline function _generic_powermeasure_stage2(μ::WeightedMeasure, exponent::Tuple)
+    ν = μ.base^exponent
+    k = maybestatic_length(ν) * μ.logweight
+    return weightedmeasure(k, ν)
 end
 
 ###############################################################################
 # ProductMeasure
 
-productmeasure(mar::FillArrays.Fill) = powermeasure(mar.value, mar.axes)
+"""
+    productmeasure(μs)
+
+Constructs a product over a collection `μs` of measures.
+
+Examples:
+
+```julia
+productmeasure((StdNormal(), StdExponential()))
+productmeasure((a = StdNormal(), b = StdExponential()))
+productmeasure([pushfwd(Base.Fix1(*, scale), StdExponential()) for scale in 0.1:0.2:2])
+```
+"""
+function productmeasure end
+export productmeasure
+
+@inline productmeasure(mar) = _generic_productmeasure_impl(mar)
+
+@inline _generic_productmeasure_impl(mar::FillArrays.Fill) =
+    powermeasure(_fill_value(mar), _fill_axes(mar))
+
+@inline _generic_productmeasure_impl(mar::Tuple{Vararg{AbstractMeasure}}) =
+    ProductMeasure(mar)
+_generic_productmeasure_impl(mar::Tuple{Vararg{Dirac}}) = Dirac(map(m -> m.x, mar))
+_generic_productmeasure_impl(mar::Tuple) = productmeasure(map(asmeasure, mar))
+
+@inline _generic_productmeasure_impl(
+    mar::NamedTuple{names,<:Tuple{Vararg{AbstractMeasure}}},
+) where {names} = ProductMeasure(mar)
+_generic_productmeasure_impl(mar::NamedTuple{names,<:Tuple{Vararg{Dirac}}}) where {names} =
+    Dirac(map(m -> m.x, mar))
+_generic_productmeasure_impl(mar::NamedTuple) = productmeasure(map(asmeasure, mar))
+
+@inline _generic_productmeasure_impl(mar::AbstractArray{<:AbstractProductMeasure}) =
+    ProductMeasure(mar)
+
+_generic_productmeasure_impl(mar::AbstractArray{<:Dirac}) = Dirac((m -> m.x).(mar))
+
+# TODO: We should be able to further optimize this
+function _generic_productmeasure_impl(mar::AbstractArray{T}) where {T}
+    if Base.issingletontype(T)
+        first(mar)^size(mar)
+    else
+        ProductMeasure(asmeasure.(mar))
+    end
+end
 
-function productmeasure(mar::ReadonlyMappedArray{T,N,A,Returns{M}}) where {T,N,A,M}
+@inline function _generic_productmeasure_impl(
+    mar::ReadonlyMappedArray{T,N,A,Returns{M}},
+) where {T,N,A,M}
     return powermeasure(mar.f.value, axes(mar.data))
 end
 
-productmeasure(mar::Base.Generator) = ProductMeasure(mar)
-productmeasure(mar::AbstractArray) = ProductMeasure(mar)
+@inline _generic_productmeasure_impl(mar::Base.Generator) = ProductMeasure(mar)
 
 # TODO: Make this static when its length is static
-@inline function productmeasure(
-    mar::AbstractArray{WeightedMeasure{StaticFloat64{W},M}},
+@inline function _generic_productmeasure_impl(
+    mar::AbstractArray{<:WeightedMeasure{StaticFloat64{W},M}},
 ) where {W,M}
     return weightedmeasure(W * length(mar), productmeasure(map(basemeasure, mar)))
 end
 
-productmeasure(nt::NamedTuple) = ProductMeasure(nt)
-productmeasure(tup::Tuple) = ProductMeasure(tup)
+# ToDo: Remove or at least refactor this (ProductMeasure shouldn't take a kernel as its argument).
 
-productmeasure(f, param_maps, pars) = ProductMeasure(kernel(f, param_maps), pars)
+productmeasure(f, param_maps, pars) = productmeasure(kernel(f, param_maps), pars)
 
 function productmeasure(k::ParameterizedTransitionKernel, pars)
     productmeasure(k.suff, k.param_maps, pars)

From fe1d136b245a3fbdd6d2d41dfcaa1aaf65209e9a Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 13:46:21 +0200
Subject: [PATCH 042/122] Generalized transport for products and unknown-DOF
 measures
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add the generalized product/power transport machinery: transport
between arbitrary measures and powers of standard measures now goes
through `transport_to_mvstd` and `transport_from_mvstd_with_rest`. The
with-rest protocol consumes only as much of the flat standard variate
as each component needs, so products (and later monadic binds) whose
components have no efficiently computable DOF become transportable via
their transport origins. Multi-dimensional power measures now flatten
to one-dimensional powers as their transport origin, one-dimensional
powers pull back to powers of their parent's origin, products over
`Fill` pull back to power measures and NamedTuple-products to
Tuple-products. Standard-measure type-based transport partner selection
(`transport_to(StdNormal, μ)`) now uses `some_dof` and lives with the
new machinery, replacing the DOF-slicing implementation in
stdmeasure.jl. Also adds `NoMSpaceElementSize`/`mspace_elsize`/
`some_mspace_elsize`, equality for `AsMeasure` wrappers, and makes
mixed static/dynamic vector views AD-friendly.

Ports the original measure-algebra commits 09b12dc, ab140a1 and related
STASH commits, with fixes: the power-measure and product-measure
transport origins gained the missing `to_origin` implementations, and
the mvstd gateway methods route through the origin machinery to avoid
infinite dispatch recursion.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/MeasureBase.jl                   |   5 +
 src/collection_utils.jl              |   2 +-
 src/combinators/product.jl           |   6 +-
 src/combinators/product_transport.jl | 409 +++++++++++++++++++++++++++
 src/combinators/reshape.jl           |  15 +-
 src/mspace.jl                        |  46 +++
 src/standard/stdmeasure.jl           | 139 ---------
 7 files changed, 465 insertions(+), 157 deletions(-)
 create mode 100644 src/combinators/product_transport.jl
 create mode 100644 src/mspace.jl

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 2aa92089..67307979 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -116,6 +116,9 @@ struct AsMeasure{T} <: AbstractMeasure
     AsMeasure{T}(obj::T) where {T} = new(obj)
 end
 
+Base.:(==)(a::AsMeasure, b::AsMeasure) = a.obj == b.obj
+Base.isapprox(a::AsMeasure, b::AsMeasure; kwargs...) = isapprox(a.obj, b.obj; kwargs...)
+
 function Pretty.quoteof(d::M) where {M<:AbstractMeasure}
     the_names = fieldnames(typeof(d))
     :($M($([getfield(d, n) for n in the_names]...)))
@@ -167,6 +170,7 @@ using IrrationalConstants: loghalf
 
 include("collection_utils.jl")
 include("smf.jl")
+include("mspace.jl")
 include("getdof.jl")
 include("transport.jl")
 include("schema.jl")
@@ -205,6 +209,7 @@ include("standard/stduniform.jl")
 include("standard/stdexponential.jl")
 include("standard/stdlogistic.jl")
 include("standard/stdnormal.jl")
+include("combinators/product_transport.jl")
 include("combinators/half.jl")
 
 #include("implicitmaps.jl")
diff --git a/src/collection_utils.jl b/src/collection_utils.jl
index 1ef9f9fd..d15131d0 100644
--- a/src/collection_utils.jl
+++ b/src/collection_utils.jl
@@ -27,7 +27,7 @@ Base.@propagate_inbounds _as_tuple(v::AbstractVector, ::Val{N}) where {N} = Tupl
 
 
 Base.@propagate_inbounds function _get_or_view(A::AbstractVector, from::IntegerLike, until::IntegerLike)
-    view(A, from:until)
+    view(A, dynamic(from):dynamic(until))
 end
 
 Base.@propagate_inbounds function _get_or_view(
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 468c38c0..656ded0d 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -28,9 +28,6 @@ Base.size(μ::AbstractProductMeasure) = size(marginals(μ))
 
 basemeasure(d::AbstractProductMeasure) = productmeasure(map(basemeasure, marginals(d)))
 
-proxy(μ::ProductMeasure{<:FillArrays.Fill}) =
-    powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ)))
-
 function Base.rand(rng::AbstractRNG, ::Type{T}, d::AbstractProductMeasure) where {T}
     mar = marginals(d)
     _rand_product(rng, T, mar, eltype(mar))
@@ -85,6 +82,9 @@ struct ProductMeasure{M} <: AbstractProductMeasure
     marginals::M
 end
 
+proxy(μ::ProductMeasure{<:FillArrays.Fill}) =
+    powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ)))
+
 @inline function logdensity_rel(μ::ProductMeasure, ν::ProductMeasure, x)
     mapreduce(logdensity_rel, +, marginals(μ), marginals(ν), x)
 end
diff --git a/src/combinators/product_transport.jl b/src/combinators/product_transport.jl
new file mode 100644
index 00000000..a2b44b1b
--- /dev/null
+++ b/src/combinators/product_transport.jl
@@ -0,0 +1,409 @@
+"""
+    transport_to(ν, ::Type{MU}) where {MU<:StdMeasure}
+    transport_to(::Type{NU}, μ) where {NU<:StdMeasure}
+
+As a user convenience, a standard measure type like [`StdUniform`](@ref),
+[`StdExponential`](@ref), [`StdNormal`](@ref) or [`StdLogistic`](@ref)
+may be used directly as the source or target of a measure transport.
+
+Depending on [`MeasureBase.some_dof(μ)`](@ref) (resp. `ν`), an instance of
+the standard measure itself or a power of it will be automatically chosen as
+the transport partner.
+
+Example:
+
+```julia
+transport_to(StdNormal, μ)
+transport_to(ν, StdNormal)
+```
+"""
+function transport_to(ν, ::Type{MU}) where {MU<:StdMeasure}
+    transport_to(ν, _std_tp_partner(MU, ν))
+end
+
+function transport_to(::Type{NU}, μ) where {NU<:StdMeasure}
+    transport_to(_std_tp_partner(NU, μ), μ)
+end
+
+function transport_to(::Type{NU}, ::Type{MU}) where {NU<:StdMeasure,MU<:StdMeasure}
+    throw(
+        ArgumentError(
+            "Can't construct a transport function between the types of two standard measures, need a measure instance on one side",
+        ),
+    )
+end
+
+_std_tp_partner(::Type{M}, μ) where {M<:StdMeasure} = _std_tp_partner_bydof(M, some_dof(μ))
+_std_tp_partner_bydof(::Type{M}, ::StaticInteger{1}) where {M<:StdMeasure} = M()
+_std_tp_partner_bydof(::Type{M}, dof::IntegerLike) where {M<:StdMeasure} = M()^dof
+function _std_tp_partner_bydof(::Type{M}, ::AbstractNoDOF{MU}) where {M<:StdMeasure,MU}
+    throw(
+        ArgumentError(
+            "Can't determine a standard transport partner for measures of type $(nameof(MU))",
+        ),
+    )
+end
+
+
+# For transport, always pull a multi-dimensional PowerMeasure back to a
+# one-dimensional PowerMeasure first:
+
+const _PowerMeasureRank1{M} = PowerMeasure{M,<:NTuple{1,OneToLike}}
+
+function transport_origin(μ::PowerMeasure)
+    pwr_base(μ)^prod(pwr_size(μ))
+end
+
+function to_origin(μ::PowerMeasure, x)
+    maybestatic_reshape(x, (prod(pwr_size(μ)),))
+end
+
+function from_origin(μ::PowerMeasure, x_origin)
+    # Sanity check, should never fail:
+    @assert x_origin isa AbstractVector
+    maybestatic_reshape(x_origin, pwr_size(μ))
+end
+
+
+# A one-dimensional PowerMeasure has an origin if its parent has an origin:
+
+function transport_origin(μ::_PowerMeasureRank1)
+    _pwr_origin(typeof(μ), transport_origin(pwr_base(μ)), pwr_axes(μ))
+end
+_pwr_origin(::Type{MU}, parent_origin, axes) where {MU} = parent_origin^axes
+_pwr_origin(::Type{MU}, ::NoTransportOrigin, axes) where {MU} = NoTransportOrigin{MU}()
+
+function to_origin(μ::_PowerMeasureRank1, x)
+    to_origin.(Ref(pwr_base(μ)), x)
+end
+
+function from_origin(μ::_PowerMeasureRank1, x_origin)
+    # Sanity check, should never fail:
+    @assert x_origin isa AbstractVector
+    from_origin.(Ref(pwr_base(μ)), x_origin)
+end
+
+
+# Transport between powers of standard measures, of any rank:
+
+function _stdpow_transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, x) where {NU,MU}
+    y = transport_to(pwr_base(ν), pwr_base(μ)).(x)
+    maybestatic_reshape(y, pwr_size(ν))
+end
+
+function _stdpow_transport_def(ν::StdPowerMeasure{MU}, μ::StdPowerMeasure{MU}, x) where {MU}
+    maybestatic_reshape(x, pwr_size(ν))
+end
+
+transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, x) where {NU,MU} =
+    _stdpow_transport_def(ν, μ, x)
+
+# Disambiguation with the mvstd gateway methods below:
+transport_def(ν::StdPowerMeasure{NU,1}, μ::StdPowerMeasure{MU}, x) where {NU,MU} =
+    _stdpow_transport_def(ν, μ, x)
+transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU,1}, x) where {NU,MU} =
+    _stdpow_transport_def(ν, μ, x)
+transport_def(ν::StdPowerMeasure{NU,1}, μ::StdPowerMeasure{MU,1}, x) where {NU,MU} =
+    _stdpow_transport_def(ν, μ, x)
+
+
+# Transport between univariate standard measures and one-dimensional power
+# measures of size one:
+
+function transport_def(ν::StdMeasure, μ::StdPowerMeasure{MU,1}, x) where {MU}
+    return transport_def(ν, pwr_base(μ), only(x))
+end
+
+function transport_def(ν::StdPowerMeasure{NU,1}, μ::StdMeasure, x) where {NU}
+    sz_ν = pwr_size(ν)
+    @assert prod(sz_ν) == 1
+    return maybestatic_fill(transport_def(pwr_base(ν), μ, x), sz_ν)
+end
+
+
+# Transport to a multivariate standard measure from any measure:
+
+function transport_def(ν::StdPowerMeasure{NU,1}, μ::AbstractMeasure, x) where {NU}
+    transport_to_mvstd(pwr_base(ν), μ, x)
+end
+
+function transport_to_mvstd(ν_inner::StdMeasure, μ::AbstractMeasure, x)
+    return _to_mvstd_withdof(ν_inner, μ, fast_dof(μ), x)
+end
+
+# For standard measures and their powers specialized `transport_def` methods
+# exist, for other measures the origin-based machinery must be used directly
+# instead of `transport_def`, to prevent infinite dispatch recursion via the
+# gateway methods above:
+const _StdOrStdPowerMeasure = Union{StdMeasure,StdPowerMeasure}
+
+_transport_def_nongateway(ν::_StdOrStdPowerMeasure, μ::_StdOrStdPowerMeasure, x) =
+    transport_def(ν, μ, x)
+function _transport_def_nongateway(ν, μ, x)
+    _transport_between_origins(ν, _origin_depth(ν), _origin_depth(μ), μ, x)
+end
+
+function _to_mvstd_withdof(ν_inner::StdMeasure, μ::AbstractMeasure, dof_μ::IntegerLike, x)
+    _transport_def_nongateway(ν_inner^dof_μ, μ, x)
+end
+
+function _to_mvstd_withdof(ν_inner::StdMeasure, μ::AbstractMeasure, ::AbstractNoDOF, x)
+    _to_mvstd_withorigin(ν_inner, μ, transport_origin(μ), x)
+end
+
+function _to_mvstd_withorigin(ν_inner::StdMeasure, μ::AbstractMeasure, μ_origin, x)
+    x_origin = to_origin(μ, x)
+    transport_to_mvstd(ν_inner, μ_origin, x_origin)
+end
+
+function _to_mvstd_withorigin(ν_inner::StdMeasure, μ::AbstractMeasure, ::NoTransportOrigin, x)
+    throw(
+        ArgumentError(
+            "Don't know how to transport values of type $(nameof(typeof(x))) from $(nameof(typeof(μ))) to a power of $(nameof(typeof(ν_inner)))",
+        ),
+    )
+end
+
+
+# Transport from a multivariate standard measure to any measure:
+
+function transport_def(ν::AbstractMeasure, μ::StdPowerMeasure{MU,1}, x) where {MU}
+    _transport_from_mvstd(ν, pwr_base(μ), x)
+end
+
+function _transport_from_mvstd(ν::AbstractMeasure, μ_inner::StdMeasure, x)
+    y, x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x)
+    if !isempty(x_rest)
+        throw(ArgumentError("Input value too long during transport"))
+    end
+    return y
+end
+
+function transport_from_mvstd_with_rest(ν::AbstractMeasure, μ_inner::StdMeasure, x)
+    dof_ν = fast_dof(ν)
+    return _from_mvstd_with_rest_withdof(ν, dof_ν, μ_inner, x)
+end
+
+function _from_mvstd_with_rest_withdof(
+    ν::AbstractMeasure,
+    dof_ν::IntegerLike,
+    μ_inner::StdMeasure,
+    x,
+)
+    len_x = maybestatic_length(x)
+
+    # Since we can't check the DOF of measures like Bind, we could "run out
+    # of x" if the original x was too short. `transport_to` below will detect
+    # this, but better throw a more informative exception here:
+    if len_x < dof_ν
+        throw(ArgumentError("Variate too short during transport"))
+    end
+
+    x_inner_dof, x_rest = _split_after(x, dof_ν)
+    y = _transport_def_nongateway(ν, μ_inner^dof_ν, x_inner_dof)
+    return y, x_rest
+end
+
+function _from_mvstd_with_rest_withdof(
+    ν::AbstractMeasure,
+    ::AbstractNoDOF,
+    μ_inner::StdMeasure,
+    x,
+)
+    _from_mvstd_with_rest_withorigin(ν, transport_origin(ν), μ_inner, x)
+end
+
+function _from_mvstd_with_rest_withorigin(
+    ν::AbstractMeasure,
+    ν_origin,
+    μ_inner::StdMeasure,
+    x,
+)
+    x_origin, x_rest = transport_from_mvstd_with_rest(ν_origin, μ_inner, x)
+    from_origin(ν, x_origin), x_rest
+end
+
+function _from_mvstd_with_rest_withorigin(
+    ν::AbstractMeasure,
+    ::NoTransportOrigin,
+    μ_inner::StdMeasure,
+    x,
+)
+    throw(
+        ArgumentError(
+            "Don't know how to transport a value of type $(nameof(typeof(x))) from a power of $(nameof(typeof(μ_inner))) to $(nameof(typeof(ν)))",
+        ),
+    )
+end
+
+
+# Transport between a standard measure and Dirac:
+
+@inline transport_from_mvstd_with_rest(ν::Dirac, ::StdMeasure, x::Any) = ν.x, x
+
+@inline transport_to_mvstd(::StdMeasure, ::Dirac, ::Any) = FillArrays.Zeros{Bool}(0)
+
+
+# Pull back from a product over a Fill to a power measure:
+
+@inline transport_origin(μ::ProductMeasure) = _marginals_tp_origin(marginals(μ))
+@inline to_origin(μ::ProductMeasure, x) = _marginals_to_origin(marginals(μ), x)
+@inline from_origin(μ::ProductMeasure, x_origin) =
+    _marginals_from_origin(marginals(μ), x_origin)
+
+_marginals_tp_origin(::Ms) where {Ms} = NoTransportOrigin{ProductMeasure{Ms}}()
+
+_marginals_tp_origin(marginals_μ::FillArrays.Fill) =
+    _fill_value(marginals_μ)^_fill_axes(marginals_μ)
+_marginals_to_origin(::FillArrays.Fill, x) = x
+_marginals_from_origin(::FillArrays.Fill, x_origin) = x_origin
+
+
+# Pull back from a NamedTuple product measure to a Tuple product measure:
+#
+# Maybe ToDo (breaking): For transport between NamedTuple-marginals we could
+# match names where possible, even if given in different order, and transport
+# between the remaining non-matching names in the order given. This may not
+# be worth the additional complexity, though, since transport is typically
+# used with a (power of a) standard measure on one side.
+
+_marginals_tp_origin(marginals_μ::NamedTuple{names}) where {names} =
+    productmeasure(values(marginals_μ))
+_marginals_to_origin(::NamedTuple{names}, x::NamedTuple{names}) where {names} = values(x)
+_marginals_from_origin(::NamedTuple{names}, x_origin::Tuple) where {names} =
+    NamedTuple{names}(x_origin)
+
+
+# Transport between two instances of ProductMeasure:
+
+transport_def(ν::ProductMeasure, μ::ProductMeasure, x) =
+    _marginal_transport_def(marginals(ν), marginals(μ), x)
+
+function _marginal_transport_def(marginals_ν, marginals_μ, x)
+    @assert size(marginals_ν) == size(marginals_μ) == size(x)  # Sanity check, should not fail
+    transport_def.(marginals_ν, marginals_μ, x)
+end
+
+function _marginal_transport_def(
+    marginals_ν::Tuple{Vararg{AbstractMeasure,N}},
+    marginals_μ::Tuple{Vararg{AbstractMeasure,N}},
+    x,
+) where {N}
+    map(transport_def, marginals_ν, marginals_μ, x)
+end
+
+function _marginal_transport_def(
+    marginals_ν::AbstractVector{<:AbstractMeasure},
+    marginals_μ::Tuple{Vararg{AbstractMeasure,N}},
+    x,
+) where {N}
+    _marginal_transport_def(_as_tuple(marginals_ν, Val(N)), marginals_μ, x)
+end
+
+function _marginal_transport_def(
+    marginals_ν::Tuple{Vararg{AbstractMeasure,N}},
+    marginals_μ::AbstractVector{<:AbstractMeasure},
+    x,
+) where {N}
+    _marginal_transport_def(marginals_ν, _as_tuple(marginals_μ, Val(N)), _as_tuple(x, Val(N)))
+end
+
+
+# Transport from a ProductMeasure to a standard measure:
+
+function transport_to_mvstd(ν_inner::StdMeasure, μ::ProductMeasure, x)
+    _marginals_to_mvstd(ν_inner, marginals(μ), x)
+end
+
+struct _TransportToMvStd{NU<:StdMeasure} <: Function end
+(::_TransportToMvStd{NU})(μ, x) where {NU} = transport_to_mvstd(NU(), μ, x)
+
+function _marginals_to_mvstd(::NU, marginals_μ::Tuple, x::Tuple) where {NU<:StdMeasure}
+    _flatten_to_rv(map(_TransportToMvStd{NU}(), marginals_μ, x))
+end
+
+function _marginals_to_mvstd(::NU, marginals_μ, x) where {NU<:StdMeasure}
+    _flatten_to_rv(broadcast(_TransportToMvStd{NU}(), marginals_μ, x))
+end
+
+
+# Transport from a standard measure to a ProductMeasure, with rest:
+
+const _MaybeUnknownDOF = Union{IntegerLike,AbstractNoDOF}
+
+const _KnownDOFs = Union{Tuple{Vararg{IntegerLike,N}} where N,StaticVector{<:IntegerLike}}
+
+function transport_from_mvstd_with_rest(ν::ProductMeasure, μ_inner::StdMeasure, x)
+    νs = marginals(ν)
+    dofs = map(fast_dof, νs)
+    return _marginals_from_mvstd_with_rest(νs, dofs, μ_inner, x)
+end
+
+function _dof_access_firstidxs(dofs::Tuple{Vararg{IntegerLike,N}}, first_idx) where {N}
+    cumsum((first_idx, dofs[begin:(end-1)]...))
+end
+
+function _dof_access_firstidxs(dofs::AbstractVector{<:IntegerLike}, first_idx)
+    # ToDo: Improve implementation (reduce memory allocations):
+    cumsum(vcat([eltype(dofs)(first_idx)], dofs[begin:(end-1)]))
+end
+
+function _split_x_by_marginals_with_rest(
+    dofs::Union{Tuple,AbstractVector},
+    x::AbstractVector{<:Real},
+)
+    x_idxs = maybestatic_eachindex(x)
+    first_idxs = _dof_access_firstidxs(dofs, maybestatic_first(x_idxs))
+    xs = map((from, n) -> _get_or_view(x, from, from + n - one(n)), first_idxs, dofs)
+    x_rest = _get_or_view(x, first_idxs[end] + dofs[end], maybestatic_last(x_idxs))
+    return xs, x_rest
+end
+
+function _marginals_from_mvstd_with_rest(
+    νs,
+    dofs::_KnownDOFs,
+    μ_inner::StdMeasure,
+    x::AbstractVector{<:Real},
+)
+    xs, x_rest = _split_x_by_marginals_with_rest(dofs, x)
+    μs = map(n -> μ_inner^n, dofs)
+    ys = map(transport_def, νs, μs, xs)
+    return ys, x_rest
+end
+
+function _marginals_from_mvstd_with_rest(
+    νs,
+    dofs,
+    μ_inner::StdMeasure,
+    x::AbstractVector{<:Real},
+)
+    _marginals_from_mvstd_with_rest_nodof(νs, μ_inner, x)
+end
+
+function _marginals_from_mvstd_with_rest_nodof(
+    νs::Tuple{Vararg{AbstractMeasure}},
+    μ_inner::StdMeasure,
+    x::AbstractVector{<:Real},
+)
+    # ToDo: Check for type stability, may need a generated function:
+    y1, x_rest = transport_from_mvstd_with_rest(νs[1], μ_inner, x)
+    y2_end, x_final_rest = _marginals_from_mvstd_with_rest_nodof(Base.tail(νs), μ_inner, x_rest)
+    return (y1, y2_end...), x_final_rest
+end
+
+_marginals_from_mvstd_with_rest_nodof(::Tuple{}, ::StdMeasure, x::AbstractVector{<:Real}) =
+    (), x
+
+function _marginals_from_mvstd_with_rest_nodof(
+    νs::AbstractVector{<:AbstractMeasure},
+    μ_inner::StdMeasure,
+    x::AbstractVector{<:Real},
+)
+    # ToDo: Check for type stability:
+    ys = Vector{Any}(undef, length(eachindex(νs)))
+    x_rest = x
+    for (i, ν) in zip(eachindex(ys), νs)
+        ys[i], x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x_rest)
+    end
+    return [y for y in ys], x_rest
+end
diff --git a/src/combinators/reshape.jl b/src/combinators/reshape.jl
index dddae55b..f77dbb68 100644
--- a/src/combinators/reshape.jl
+++ b/src/combinators/reshape.jl
@@ -53,17 +53,4 @@ a space of arrays with shape `sz`.
 function mreshape end
 
 mreshape(m::AbstractMeasure, sz::IntegerLike...) = mreshape(m, sz)
-mreshape(m::AbstractMeasure, sz::SizeLike) = pushfwd(Reshape(sz, mspace_elsize(m)), m)
-
-
-"""
-    MeasureBase.mspace_elsize(m::AbstractMeasure)::MeasureBase.SizeLike
-
-Return the size of the elements of the measurable space of `m`.
-
-Defaults to the size of a test value of `m`, may be specialized for
-measure types where this is inefficient.
-"""
-function mspace_elsize end
-
-mspace_elsize(m::AbstractMeasure) = maybestatic_size(testvalue(m))
+mreshape(m::AbstractMeasure, sz::SizeLike) = pushfwd(Reshape(sz, some_mspace_elsize(m)), m)
diff --git a/src/mspace.jl b/src/mspace.jl
new file mode 100644
index 00000000..6227537d
--- /dev/null
+++ b/src/mspace.jl
@@ -0,0 +1,46 @@
+"""
+    MeasureBase.NoMSpaceElementSize{MU}
+
+Indicates that either the measurable space of measures of type `MU` is not
+a space over arrays, or that the size of the arrays is not fixed or can not
+be easily/efficiently determined.
+"""
+struct NoMSpaceElementSize{MU} end
+
+
+"""
+    mspace_elsize(μ)
+
+For a measure `μ` over an array-valued measurable space, return the size of
+the arrays that are the elements of the space.
+
+May return [`NoMSpaceElementSize{typeof(μ)}()`](@ref).
+"""
+function mspace_elsize end
+export mspace_elsize
+
+@inline mspace_elsize(μ::AbstractMeasure) = NoMSpaceElementSize{typeof(μ)}()
+
+
+"""
+    MeasureBase.some_mspace_elsize(μ::AbstractMeasure)
+
+For a measure `μ` over an array-valued measurable space, return the size of
+an arbitrary element of the space.
+
+Use with caution, the space of some measures is made up of arrays of
+different sizes!
+
+In general, use [`mspace_elsize(μ)`](@ref) instead. `some_mspace_elsize` is
+useful if the measurable space is expected to contain only arrays of the
+same size but there is no way to prove this automatically. Algorithms that
+use the returned size should always check that it matches the size of each
+point of the space that is processed.
+"""
+function some_mspace_elsize end
+
+@inline some_mspace_elsize(μ) = _mspace_some_elsize_impl(μ, mspace_elsize(μ))
+
+@inline _mspace_some_elsize_impl(::AbstractMeasure, sz::SizeLike) = sz
+_mspace_some_elsize_impl(μ::AbstractMeasure, ::NoMSpaceElementSize) =
+    maybestatic_size(testvalue(μ))
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index 3dd30e24..81409796 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -14,142 +14,3 @@ const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}}
 @inline check_dof(::StdMeasure, ::StdMeasure) = nothing
 
 @inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x
-
-function transport_def(ν::StdMeasure, μ::PowerMeasure{<:StdMeasure}, x)
-    return transport_def(ν, μ.parent, only(x))
-end
-
-function transport_def(ν::PowerMeasure{<:StdMeasure}, μ::StdMeasure, x)
-    return maybestatic_fill(transport_def(ν.parent, μ, only(x)), map(length, ν.axes))
-end
-
-function transport_def(
-    ν::PowerMeasure{<:StdMeasure,<:NTuple{1,Base.OneTo}},
-    μ::PowerMeasure{<:StdMeasure,<:NTuple{1,Base.OneTo}},
-    x,
-)
-    return transport_to(ν.parent, μ.parent).(x)
-end
-
-function transport_def(
-    ν::PowerMeasure{<:StdMeasure,<:NTuple{N,Base.OneTo}},
-    μ::PowerMeasure{<:StdMeasure,<:NTuple{M,Base.OneTo}},
-    x,
-) where {N,M}
-    return reshape(transport_to(ν.parent, μ.parent).(x), map(length, ν.axes)...)
-end
-
-# Implement transport_to(NU::Type{<:StdMeasure}, μ) and transport_to(ν, MU::Type{<:StdMeasure}):
-
-_std_measure(::Type{M}, ::StaticInteger{1}) where {M<:StdMeasure} = M()
-_std_measure(::Type{M}, dof::IntegerLike) where {M<:StdMeasure} = M()^dof
-_std_measure_for(::Type{M}, μ::Any) where {M<:StdMeasure} = _std_measure(M, getdof(μ))
-
-function transport_to(::Type{NU}, μ) where {NU<:StdMeasure}
-    transport_to(_std_measure_for(NU, μ), μ)
-end
-
-function transport_to(ν, ::Type{MU}) where {MU<:StdMeasure}
-    transport_to(ν, _std_measure_for(MU, ν))
-end
-
-# Transform between standard measures and Dirac:
-
-@inline transport_def(ν::Dirac, ::PowerMeasure{<:StdMeasure}, ::Any) = ν.x
-
-@inline function transport_def(ν::PowerMeasure{<:StdMeasure}, ::Dirac, ::Any)
-    Zeros{Bool}(map(_ -> 0, ν.axes))
-end
-
-# Helpers for product transforms and similar:
-
-struct _TransportToStd{NU<:StdMeasure} <: Function end
-(::_TransportToStd{NU})(μ, x) where {NU} = transport_to(NU()^getdof(μ), μ)(x)
-
-struct _TransportFromStd{MU<:StdMeasure} <: Function end
-_TransportFromStd{MU}(ν, x) where {MU} = transport_to(ν, MU()^getdof(ν))(x)
-
-function _tuple_transport_def(
-    ν::PowerMeasure{NU},
-    μs::Tuple,
-    xs::Tuple,
-) where {NU<:StdMeasure}
-    reshape(vcat(map(_TransportToStd{NU}(), μs, xs)...), ν.axes)
-end
-
-function transport_def(
-    ν::PowerMeasure{NU},
-    μ::ProductMeasure{<:Tuple},
-    x,
-) where {NU<:StdMeasure}
-    _tuple_transport_def(ν, marginals(μ), x)
-end
-
-function transport_def(
-    ν::PowerMeasure{NU},
-    μ::ProductMeasure{<:NamedTuple{names}},
-    x,
-) where {NU<:StdMeasure,names}
-    _tuple_transport_def(ν, values(marginals(μ)), values(x))
-end
-
-@inline _offset_cumsum(s, x, y, rest...) = (s, _offset_cumsum(s + x, y, rest...)...)
-@inline _offset_cumsum(s, x) = (s,)
-@inline _offset_cumsum(s) = ()
-
-function _stdvar_viewranges(μs::Tuple, startidx::IntegerLike)
-    N = map(getdof, μs)
-    offs = _offset_cumsum(startidx, N...)
-    map((o, n) -> o:(o+n-1), offs, N)
-end
-
-function _tuple_transport_def(
-    νs::Tuple,
-    μ::PowerMeasure{MU},
-    x::AbstractArray{<:Real},
-) where {MU<:StdMeasure}
-    vrs = _stdvar_viewranges(νs, firstindex(x))
-    xs = map(r -> view(x, r), vrs)
-    map(_TransportFromStd{MU}, νs, xs)
-end
-
-function transport_def(
-    ν::ProductMeasure{<:Tuple},
-    μ::PowerMeasure{MU},
-    x,
-) where {MU<:StdMeasure}
-    _tuple_transport_def(marginals(ν), μ, x)
-end
-
-function transport_def(
-    ν::ProductMeasure{<:NamedTuple{names}},
-    μ::PowerMeasure{MU},
-    x,
-) where {MU<:StdMeasure,names}
-    NamedTuple{names}(_tuple_transport_def(values(marginals(ν)), μ, x))
-end
-
-function transport_def(
-    ν::PowerMeasure{NU},
-    μ::ProductMeasure{<:AbstractArray},
-    x,
-) where {NU<:StdMeasure}
-    reshape(vcat(map(_TransportToStd{NU}(), marginals(μ), x)...), ν.axes)
-end
-
-function _marginal_viewranges(μs::AbstractArray, startidx::IntegerLike)
-    ns = map(m -> dynamic(getdof(m)), μs)
-    offs = cumsum(vcat(dynamic(startidx), ns[begin:(end-1)]))
-    map((o, n) -> o:(o+n-1), offs, ns)
-end
-
-function transport_def(
-    ν::ProductMeasure{<:AbstractArray},
-    μ::PowerMeasure{MU},
-    x::AbstractArray{<:Real},
-) where {MU<:StdMeasure}
-    νs = marginals(ν)
-    vrs = _marginal_viewranges(νs, firstindex(x))
-    xs = map(r -> view(x, r), vrs)
-    map(_TransportFromStd{MU}, νs, xs)
-end

From 8c3d680a2434835b875f7817b411af27581890a1 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 15:19:40 +0200
Subject: [PATCH 043/122] Rework mbind as combined-value bind, add mcombine
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Rework the monadic bind into a full hierarchical measure: `mbind(f_β,
α, f_c)` combines points of the primary measure `α` and the dependent
secondary measure `f_β(a)` via a value combination function `f_c`
(`tuple`, `Pair`, `vcat` and `merge` are supported for density
evaluation, `OneTwoMany.secondarg` gives the plain monadic bind).
`mkernel(f_β, f_c)` represents the generalized transition kernel,
`bindkernel` and `boundmeasure` recover the parts of a bind. Density
evaluation goes through the new `transportmeasure`/
`tpmeasure_split_combined` mechanism, and transport to and from powers
of standard measures uses the with-rest protocol, so hierarchical
measures without a fast DOF are transportable. Add `mcombine` and
`CombinedMeasure` for combining two independent measures via `f_c`,
with product shortcuts for `tuple`, `vcat` and `merge`. MeasureBase now
depends on OneTwoMany for `firstarg`/`secondarg`.

Ports the original measure-algebra commits fe28297, 03724b7, 17e30ec,
cce5b41, 8561e26, 4b3ed53 and related STASH commits, with fixes:
`boundmeasure` returns the bound measure instead of the kernel, the
tuple/Pair variate splitters no longer reference undefined variables,
kernel results are converted via `asmeasure`, and NamedTuple products
gained the missing to/from-mvstd transport paths.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 Project.toml                         |   2 +
 src/MeasureBase.jl                   |   5 +-
 src/combinators/bind.jl              | 336 +++++++++++++++++++++++++--
 src/combinators/combined.jl          | 174 ++++++++++++++
 src/combinators/product_transport.jl |  18 ++
 test/Project.toml                    |   1 +
 test/combinators/bind.jl             |  94 ++++++++
 test/combinators/combined.jl         |  52 +++++
 test/runtests.jl                     |   2 +
 9 files changed, 658 insertions(+), 26 deletions(-)
 create mode 100644 src/combinators/combined.jl
 create mode 100644 test/combinators/bind.jl
 create mode 100644 test/combinators/combined.jl

diff --git a/Project.toml b/Project.toml
index e2bfee4a..c0ba40c8 100644
--- a/Project.toml
+++ b/Project.toml
@@ -23,6 +23,7 @@ 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"
@@ -81,6 +82,7 @@ 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"
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 67307979..9f8037eb 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -63,6 +63,8 @@ using HeterogeneousComputing: real_numtype
 using ArraysOfArrays:
     VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays, VectorOfSimilarVectors, flatview
 
+using OneTwoMany: firstarg, secondarg
+
 export gentype
 
 export AbstractMeasure
@@ -190,7 +192,6 @@ include("primitives/lebesgue.jl")
 include("primitives/dirac.jl")
 include("primitives/trivial.jl")
 
-include("combinators/bind.jl")
 include("combinators/transformedmeasure.jl")
 include("combinators/reshape.jl")
 include("combinators/weighted.jl")
@@ -210,6 +211,8 @@ include("standard/stdexponential.jl")
 include("standard/stdlogistic.jl")
 include("standard/stdnormal.jl")
 include("combinators/product_transport.jl")
+include("combinators/combined.jl")
+include("combinators/bind.jl")
 include("combinators/half.jl")
 
 #include("implicitmaps.jl")
diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl
index 60b7cfd8..27cb02a4 100644
--- a/src/combinators/bind.jl
+++ b/src/combinators/bind.jl
@@ -1,43 +1,329 @@
+@doc raw"""
+    mkernel(f_β, f_c = OneTwoMany.secondarg)::Function
+
+Constructs a generalized monadic transition kernel from a primary transition
+kernel function `f_β` and a value combination function `f_c`.
+
+`f_β` must behave like `β = f_β(a)`, taking a value `a` from a primary
+measurable space and returning a measure-like object `β`.
+
+`f_c` must behave like `c = f_c(a, b)`, taking a value `a` (like `f_β`) and
+a value `b` from the measurable space of `β` and returning a value `c`.
+
+`f_k = mkernel(f_β, f_c)` then acts like
+
+```julia
+f_k(a) ≡ pushfwd(c -> f_c(c[1], c[2]), productmeasure((Dirac(a), f_β(a))))
+```
+
+(`≡` denoting pseudocode-equivalency here). So with the default
+`f_c == OneTwoMany.secondarg`, we just have `f_k(a) ≡ f_β(a)`.
+
+Also,
+
+```julia
+mbind(mkernel(f_β, f_c), α) == mbind(f_β, α, f_c)
+```
+
+See also [`mbind`](@ref).
 """
-    struct MeasureBase.Bind{M,K} <: AbstractMeasure
+function mkernel end
+export mkernel
+
 
-Represents a monatic bind. User code should not create instances of `Bind`
-directly, but should call `mbind(k, μ)` instead.
 """
-struct Bind{M,K} <: AbstractMeasure
-    k::K
-    μ::M
-end
+    struct MeasureBase.MKernel <: Function
 
-getdof(d::Bind) = NoDOF{typeof(d)}()
+Represents a generalized monadic transition kernel.
 
-function Base.rand(rng::AbstractRNG, ::Type{T}, d::Bind) where {T}
-    x = rand(rng, T, d.μ)
-    y = rand(rng, T, d.k(x))
-    return y
+User code should not create instances of `MKernel` directly, but should
+call [`mkernel`](@ref) instead.
+"""
+struct MKernel{FT,FC} <: Function
+    f_β::FT
+    f_c::FC
 end
 
+(f_k::MKernel)(a) = mbind(f_k, Dirac(a))
 
-"""
-    mbind(k, μ)::AbstractMeasure
+@inline mkernel(f_β::MKernel) = f_β
+@inline mkernel(f_β, f_c = secondarg) = _generic_mkernel_impl(f_β, f_c)
+
+@inline _generic_mkernel_impl(f_β, f_c) = MKernel(f_β, f_c)
+@inline _generic_mkernel_impl(f_β::MKernel, ::typeof(secondarg)) = f_β
+
+
+@doc raw"""
+    mbind(f_β, α::AbstractMeasure, f_c = OneTwoMany.secondarg)
+    mbind(f_β::MeasureBase.MKernel, α::AbstractMeasure)
+
+Constructs a monadic bind, resp. a hierarchical measure, from a transition
+kernel function `f_β`, a primary measure `α` and a value combination
+function `f_c`.
+
+`f_β` must be a function that maps a point `a` from the space of the primary
+measure `α` to a dependent secondary measure `β_a = f_β(a)`.
+`ab = f_c(a, b)` must map such a point `a` and a point `b` from the
+space of measure `β_a` to a combined value `ab = f_c(a, b)`.
+
+The resulting measure
 
-Given
+```julia
+μ = mbind(f_β, α, f_c)
+```
 
-- a measure μ
-- a kernel function k that takes values from the support of μ and returns a
-  measure
+has the mathematical interpretation (on sets $$A$$ and $$B$$)
 
-The *monadic bind* operation `mbind(k, μ)` returns is a new measure.
+```math
+\mu(f_c(A, B)) = \int_A \beta_a(B)\, \mathrm{d}\, \alpha(a)
+```
 
-A monadic bind is often written as `>>=` (e.g. in Haskell), but this symbol is
-unavailable in Julia.
+When using the default `f_c = OneTwoMany.secondarg` (so `ab == b`) this
+simplifies to
 
+```math
+\mu(B) = \int_A \beta_a(B)\, \mathrm{d}\, \alpha(a)
 ```
-μ = StdExponential()
-ν = mbind(μ) do scale
-    pushfwd(Base.Fix1(*, scale), StdNormal())
+
+which is equivalent to a monadic bind, viewing measures as monads.
+
+Computationally, `ab = rand(μ)` is equivalent to
+
+```julia
+a = rand(α)
+β_a = f_β(a)
+b = rand(β_a)
+ab = f_c(a, b)
+```
+
+The measure `α` that went into the bind can be retrieved via
+`boundmeasure(mbind(f_β, α, f_c)) == α` and the kernel via
+`bindkernel(mbind(f_β, α, f_c)) == mkernel(f_β, f_c)`.
+
+Densities on hierarchical measures can only be evaluated if `ab = f_c(a, b)`
+can be unambiguously split into `a` and `b` again, knowing `α`. This is
+currently implemented for `f_c` that is either `tuple` or `=>`/`Pair` (these
+work for any combination of variate types), `vcat` (for tuple- or
+vector-like variates) and `merge` (`NamedTuple` variates).
+[`MeasureBase.tpmeasure_split_combined`](@ref) can be specialized to
+support other choices for `f_c`.
+
+# Extended help
+
+Bayesian example with a correlated prior: Mathematically, let
+
+    position = a1 ~ StdNormal()
+    noise = a2 ~ pushforward(h(a1, ·), StdExponential())
+
+where `h(a1, a2) = √(abs(a1) * a2)`. Because this prior on the space of
+`A = A1 × A2 = (position, noise)` is a hierarchical measure (`a2` depends
+on `a1`), we can construct it using `mbind` with `merge` as `f_c`:
+
+```julia
+using MeasureBase, AffineMaps
+
+prior = mbind(
+    productmeasure((
+        position = StdNormal(),
+    )), merge
+) do a
+    productmeasure((
+        noise = pushfwd(setinverse(sqrt, setladj(x -> x^2, x -> log(2))) ∘ Mul(abs(a.position)), StdExponential()),
+    ))
 end
+
+model = θ -> pushfwd(MulAdd(θ.noise, θ.position), StdNormal())^10
+
+joint_θ_obs = mbind(model, prior, tuple)
+prior_predictive = mbind(model, prior)
+
+observation = rand(prior_predictive)
+likelihood = likelihoodof(model, observation)
+
+posterior = mintegrate(likelihood, prior)
+
+θ = rand(prior)
+logdensityof(posterior, θ)
 ```
 """
-mbind(k, μ) = Bind(k, μ)
+function mbind end
 export mbind
+
+@inline mbind(f_β) = Base.Fix1(mbind, f_β)
+
+@inline function mbind(f_β, α::AbstractMeasure, f_c = secondarg)
+    _generic_mbind_impl(f_β, asmeasure(α), f_c)
+end
+
+@inline function _generic_mbind_impl(f_β, α::AbstractMeasure, f_c)
+    F, M, G = Core.Typeof(f_β), Core.Typeof(α), Core.Typeof(f_c)
+    Bind{F,M,G}(f_β, α, f_c)
+end
+
+@inline _generic_mbind_impl(f_β, α::Dirac, f_c) = mcombine(f_c, α, asmeasure(f_β(α.x)))
+
+@inline _generic_mbind_impl(@nospecialize(f_β), α::AbstractMeasure, ::typeof(firstarg)) = α
+@inline _generic_mbind_impl(@nospecialize(f_β), α::Dirac, ::typeof(firstarg)) = α
+
+@inline _generic_mbind_impl(f_k::MKernel, α::AbstractMeasure, ::typeof(secondarg)) =
+    mbind(f_k.f_β, α, f_k.f_c)
+@inline _generic_mbind_impl(f_k::MKernel, α::Dirac, ::typeof(secondarg)) =
+    mbind(f_k.f_β, α, f_k.f_c)
+
+
+"""
+    struct MeasureBase.Bind <: AbstractMeasure
+
+Represents a monadic bind resp. a hierarchical measure in general.
+
+User code should not create instances of `Bind` directly, but should call
+[`mbind`](@ref) instead.
+"""
+struct Bind{FT,M<:AbstractMeasure,FC} <: AbstractMeasure
+    f_β::FT
+    α::M
+    f_c::FC
+end
+
+# ToDo: Store MKernel in Bind instead of separate fields f_β and f_c?
+
+
+"""
+    bindkernel(μ::Bind)::MKernel
+
+Returns the monadic transition kernel of a monadic bind, so that
+`bindkernel(mbind(f_k::MKernel, α)) == f_k`.
+
+See [`mbind`](@ref) and [`mkernel`](@ref) for details.
+"""
+function bindkernel end
+export bindkernel
+
+bindkernel(μ::Bind) = mkernel(μ.f_β, μ.f_c)
+
+
+"""
+    boundmeasure(μ::Bind)::AbstractMeasure
+
+Returns the measure that went into a monadic bind, so that
+`boundmeasure(mbind(f_k, α)) == α`.
+
+See [`mbind`](@ref) and [`mkernel`](@ref) for details.
+"""
+function boundmeasure end
+export boundmeasure
+
+boundmeasure(μ::Bind) = μ.α
+
+
+_get_β_a(μ::Bind, a) = asmeasure(μ.f_β(a))
+
+function transportmeasure(μ::Bind, x)
+    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, x)
+    tpm_β_a = transportmeasure(_get_β_a(μ, a), b)
+    mcombine(μ.f_c, tpm_α, tpm_β_a)
+end
+
+localmeasure(μ::Bind, x) = transportmeasure(μ, x)
+
+tpmeasure_split_combined(f_c, μ::Bind, xy) = _bind_tpm_sc(f_c, μ, xy)
+
+function _bind_tpm_sc(::typeof(tuple), μ::Bind, xy::Tuple{Vararg{Any,2}})
+    x, y = xy[1], xy[2]
+    tpm_μ = transportmeasure(μ, x)
+    return tpm_μ, x, y
+end
+
+function _bind_tpm_sc(::Type{Pair}, μ::Bind, xy::Pair)
+    x, y = xy.first, xy.second
+    tpm_μ = transportmeasure(μ, x)
+    return tpm_μ, x, y
+end
+
+const _BindBy{FC} = Bind{<:Any,<:AbstractMeasure,FC}
+_bind_tpm_sc(f_c::typeof(vcat), μ::_BindBy{typeof(vcat)}, xy::AbstractVector) =
+    _bind_tpm_sc_cat(f_c, μ, xy)
+_bind_tpm_sc(f_c::typeof(merge), μ::_BindBy{typeof(merge)}, xy::NamedTuple) =
+    _bind_tpm_sc_cat(f_c, μ, xy)
+
+function _bind_tpm_sc_cat_lμabyxy(f_c, μ, xy)
+    tpm_α, a, by = tpmeasure_split_combined(μ.f_c, μ.α, xy)
+    β_a = _get_β_a(μ, a)
+    tpm_β_a, b, y = tpmeasure_split_combined(f_c, β_a, by)
+    tpm_μ = mcombine(μ.f_c, tpm_α, tpm_β_a)
+    return tpm_μ, a, b, y, xy
+end
+
+function _bind_tpm_sc_cat(f_c::typeof(vcat), μ::_BindBy{typeof(vcat)}, xy::AbstractVector)
+    tpm_μ, a, b, y, xy = _bind_tpm_sc_cat_lμabyxy(f_c, μ, xy)
+    # Don't use `x = f_c(a, b)` here, would allocate, splitting xy can use views:
+    x, y = _split_after(xy, length(a) + length(b))
+    return tpm_μ, x, y
+end
+
+function _bind_tpm_sc_cat(f_c::typeof(merge), μ::_BindBy{typeof(merge)}, xy::NamedTuple)
+    tpm_μ, a, b, y, xy = _bind_tpm_sc_cat_lμabyxy(f_c, μ, xy)
+    return tpm_μ, f_c(a, b), y
+end
+
+
+@inline insupport(μ::Bind, ::Any) = NoFastInsupport{typeof(μ)}()
+
+@inline getdof(μ::Bind) = NoDOF{typeof(μ)}()
+
+# Bypass `checked_arg`, would require potentially costly evaluation of f_β:
+@inline checked_arg(::Bind, x) = x
+
+rootmeasure(::Bind) =
+    throw(ArgumentError("root measure is implicit, but can't be instantiated, for Bind"))
+
+basemeasure(::Bind) = throw(ArgumentError("basemeasure is not available for Bind"))
+
+testvalue(::Bind) = throw(ArgumentError("testvalue is not available for Bind"))
+
+logdensity_def(::Bind, x) =
+    throw(ArgumentError("logdensity_def is not available for Bind"))
+
+# Specialize logdensityof to avoid duplicate calculations:
+function logdensityof(μ::Bind, x)
+    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, x)
+    β_a = _get_β_a(μ, a)
+    logdensityof(tpm_α, a) + logdensityof(β_a, b)
+end
+
+# Specialize unsafe_logdensityof to avoid duplicate calculations:
+function unsafe_logdensityof(μ::Bind, x)
+    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, x)
+    β_a = _get_β_a(μ, a)
+    unsafe_logdensityof(tpm_α, a) + unsafe_logdensityof(β_a, b)
+end
+
+
+function Base.rand(rng::Random.AbstractRNG, ::Type{T}, μ::Bind) where {T<:Real}
+    a = rand(rng, T, μ.α)
+    b = rand(rng, T, _get_β_a(μ, a))
+    return μ.f_c(a, b)
+end
+
+function Base.rand(rng::Random.AbstractRNG, μ::Bind)
+    a = rand(rng, μ.α)
+    b = rand(rng, _get_β_a(μ, a))
+    return μ.f_c(a, b)
+end
+
+
+function transport_to_mvstd(ν_inner::StdMeasure, μ::Bind, ab)
+    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab)
+    β_a = _get_β_a(μ, a)
+    y1 = transport_to_mvstd(ν_inner, tpm_α, a)
+    y2 = transport_to_mvstd(ν_inner, β_a, b)
+    return vcat(y1, y2)
+end
+
+
+function transport_from_mvstd_with_rest(ν::Bind, μ_inner::StdMeasure, x)
+    a, x2 = transport_from_mvstd_with_rest(ν.α, μ_inner, x)
+    β_a = _get_β_a(ν, a)
+    b, x_rest = transport_from_mvstd_with_rest(β_a, μ_inner, x2)
+    return ν.f_c(a, b), x_rest
+end
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
new file mode 100644
index 00000000..c3ea079c
--- /dev/null
+++ b/src/combinators/combined.jl
@@ -0,0 +1,174 @@
+"""
+    MeasureBase.tpmeasure_split_combined(f_c, α::AbstractMeasure, ab)
+
+Splits a combined value `ab` that originated from combining a point `a`
+from the space of a measure `α` with a point `b` from the space of
+another measure `β` via `ab = f_c(a, b)`.
+
+Returns a semantic equivalent of
+`(MeasureBase.transportmeasure(α, a), a, b)`.
+
+With `a_orig = rand(α)`, `b_orig = rand(β)` and
+`ab = f_c(a_orig, b_orig)`, the following must hold true:
+
+```julia
+tpm_α, a, b = tpmeasure_split_combined(f_c, α, ab)
+a ≈ a_orig && b ≈ b_orig
+```
+"""
+function tpmeasure_split_combined end
+
+function tpmeasure_split_combined(f_c, α::AbstractMeasure, ab)
+    a, b = _generic_split_combined(f_c, α, ab)
+    return transportmeasure(α, a), a, b
+end
+
+@inline _generic_split_combined(::typeof(tuple), ::AbstractMeasure, x::Tuple{Vararg{Any,2}}) = x
+@inline _generic_split_combined(::Type{Pair}, ::AbstractMeasure, ab::Pair) = (ab...,)
+
+function _generic_split_combined(f_c::FC, α::AbstractMeasure, ab) where {FC}
+    _split_variate_byvalue(f_c, testvalue(α), ab)
+end
+
+_split_variate_byvalue(::typeof(vcat), test_a::AbstractVector, ab::AbstractVector) =
+    _split_after(ab, length(test_a))
+
+_split_variate_byvalue(::typeof(vcat), ::NTuple{N,Any}, ab::Tuple) where {N} =
+    _split_after(ab, Val{N}())
+
+function _split_variate_byvalue(::typeof(merge), ::NamedTuple{names_a}, ab::NamedTuple) where {names_a}
+    _split_after(ab, Val(names_a))
+end
+
+
+@doc raw"""
+    mcombine(f_c, α::AbstractMeasure, β::AbstractMeasure)
+
+Combines two measures `α` and `β` to a combined measure via a point
+combination function `f_c`.
+
+`f_c` must combine a given point `a` from the space of measure `α` with a
+given point `b` from the space of measure `β` to a single value
+`ab = f_c(a, b)` in the space of the combined measure
+`μ = mcombine(f_c, α, β)`.
+
+The combined measure has the mathematical interpretation (on sets
+$$A$$ and $$B$$)
+
+```math
+\mu(f_c(A, B)) = \alpha(A)\, \beta(B)
+```
+"""
+function mcombine end
+export mcombine
+
+@inline function mcombine(f_c, α::AbstractMeasure, β::AbstractMeasure)
+    _generic_mcombine_impl_stage1(f_c, α, β)
+end
+
+@inline _generic_mcombine_impl_stage1(::typeof(firstarg), α::AbstractMeasure, β::AbstractMeasure) = α
+@inline _generic_mcombine_impl_stage1(::typeof(secondarg), α::AbstractMeasure, β::AbstractMeasure) = β
+
+@inline function _generic_mcombine_impl_stage1(::typeof(tuple), α::AbstractMeasure, β::AbstractMeasure)
+    productmeasure((α, β))
+end
+
+@inline function _generic_mcombine_impl_stage1(
+    f_c::Union{typeof(vcat),typeof(merge)},
+    α::AbstractProductMeasure,
+    β::AbstractProductMeasure,
+)
+    _mcombine_product_shortcut(f_c, marginals(α), marginals(β), α, β)
+end
+
+_mcombine_product_shortcut(::typeof(vcat), ma::AbstractVector, mb::AbstractVector, α, β) =
+    productmeasure(vcat(ma, mb))
+_mcombine_product_shortcut(::typeof(merge), ma::NamedTuple, mb::NamedTuple, α, β) =
+    productmeasure(merge(ma, mb))
+_mcombine_product_shortcut(f_c, ma, mb, α, β) = _generic_mcombine_impl_stage2(f_c, α, β)
+
+@inline function _generic_mcombine_impl_stage1(f_c, α::AbstractMeasure, β::AbstractMeasure)
+    _generic_mcombine_impl_stage2(f_c, α, β)
+end
+
+@inline function _generic_mcombine_impl_stage2(f_c, α::AbstractMeasure, β::AbstractMeasure)
+    FC, MA, MB = Core.Typeof(f_c), Core.Typeof(α), Core.Typeof(β)
+    CombinedMeasure{FC,MA,MB}(f_c, α, β)
+end
+
+@inline function _generic_mcombine_impl_stage2(f_c, α::Dirac, β::Dirac)
+    Dirac(f_c(α.x, β.x))
+end
+
+
+"""
+    struct CombinedMeasure <: AbstractMeasure
+
+Represents a combination of two measures.
+
+User code should not create instances of `CombinedMeasure` directly, but
+should call [`mcombine(f_c, α, β)`](@ref) instead.
+"""
+struct CombinedMeasure{FC,MA<:AbstractMeasure,MB<:AbstractMeasure} <: AbstractMeasure
+    f_c::FC
+    α::MA
+    β::MB
+end
+
+
+@inline insupport(μ::CombinedMeasure, ab) = NoFastInsupport{typeof(μ)}()
+
+@inline getdof(μ::CombinedMeasure) = getdof(μ.α) + getdof(μ.β)
+@inline fast_dof(μ::CombinedMeasure) = fast_dof(μ.α) + fast_dof(μ.β)
+
+# Bypass `checked_arg`, would require splitting ab:
+@inline checked_arg(::CombinedMeasure, ab) = ab
+
+rootmeasure(μ::CombinedMeasure) = mcombine(μ.f_c, rootmeasure(μ.α), rootmeasure(μ.β))
+
+basemeasure(μ::CombinedMeasure) = mcombine(μ.f_c, basemeasure(μ.α), basemeasure(μ.β))
+
+function logdensity_def(μ::CombinedMeasure, ab)
+    # Use tpmeasure_split_combined to avoid duplicate calculation of transportmeasure(α):
+    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab)
+    return logdensity_def(tpm_α, a) + logdensity_def(μ.β, b)
+end
+
+# Specialize logdensityof directly to avoid creating temporary combined base measures:
+function logdensityof(μ::CombinedMeasure, ab)
+    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab)
+    return logdensityof(tpm_α, a) + logdensityof(μ.β, b)
+end
+
+function unsafe_logdensityof(μ::CombinedMeasure, ab)
+    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab)
+    return unsafe_logdensityof(tpm_α, a) + unsafe_logdensityof(μ.β, b)
+end
+
+
+function Base.rand(rng::Random.AbstractRNG, ::Type{T}, μ::CombinedMeasure) where {T<:Real}
+    a = rand(rng, T, μ.α)
+    b = rand(rng, T, μ.β)
+    return μ.f_c(a, b)
+end
+
+function Base.rand(rng::Random.AbstractRNG, μ::CombinedMeasure)
+    a = rand(rng, μ.α)
+    b = rand(rng, μ.β)
+    return μ.f_c(a, b)
+end
+
+
+function transport_to_mvstd(ν_inner::StdMeasure, μ::CombinedMeasure, ab)
+    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab)
+    y1 = transport_to_mvstd(ν_inner, tpm_α, a)
+    y2 = transport_to_mvstd(ν_inner, μ.β, b)
+    return vcat(y1, y2)
+end
+
+
+function transport_from_mvstd_with_rest(ν::CombinedMeasure, μ_inner::StdMeasure, x)
+    a, x2 = transport_from_mvstd_with_rest(ν.α, μ_inner, x)
+    b, x_rest = transport_from_mvstd_with_rest(ν.β, μ_inner, x2)
+    return ν.f_c(a, b), x_rest
+end
diff --git a/src/combinators/product_transport.jl b/src/combinators/product_transport.jl
index a2b44b1b..01c1016c 100644
--- a/src/combinators/product_transport.jl
+++ b/src/combinators/product_transport.jl
@@ -322,6 +322,14 @@ function _marginals_to_mvstd(::NU, marginals_μ::Tuple, x::Tuple) where {NU<:Std
     _flatten_to_rv(map(_TransportToMvStd{NU}(), marginals_μ, x))
 end
 
+function _marginals_to_mvstd(
+    ν::NU,
+    marginals_μ::NamedTuple{names},
+    x::NamedTuple{names},
+) where {NU<:StdMeasure,names}
+    _marginals_to_mvstd(ν, values(marginals_μ), values(x))
+end
+
 function _marginals_to_mvstd(::NU, marginals_μ, x) where {NU<:StdMeasure}
     _flatten_to_rv(broadcast(_TransportToMvStd{NU}(), marginals_μ, x))
 end
@@ -339,6 +347,16 @@ function transport_from_mvstd_with_rest(ν::ProductMeasure, μ_inner::StdMeasure
     return _marginals_from_mvstd_with_rest(νs, dofs, μ_inner, x)
 end
 
+function transport_from_mvstd_with_rest(
+    ν::ProductMeasure{<:NamedTuple{names}},
+    μ_inner::StdMeasure,
+    x,
+) where {names}
+    ys, x_rest =
+        transport_from_mvstd_with_rest(productmeasure(values(marginals(ν))), μ_inner, x)
+    return NamedTuple{names}(ys), x_rest
+end
+
 function _dof_access_firstidxs(dofs::Tuple{Vararg{IntegerLike,N}}, first_idx) where {N}
     cumsum((first_idx, dofs[begin:(end-1)]...))
 end
diff --git a/test/Project.toml b/test/Project.toml
index cb68fc93..32833feb 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -17,6 +17,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+OneTwoMany = "762dc654-8631-413a-a342-372a7419ad9d"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 PropertyFunctions = "09e99361-2bb8-48a2-a80f-de58f0739eb4"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl
new file mode 100644
index 00000000..a1d8865a
--- /dev/null
+++ b/test/combinators/bind.jl
@@ -0,0 +1,94 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using Random
+using StableRNGs: StableRNG
+using AffineMaps: Mul
+
+using MeasureBase
+using MeasureBase: StdExponential, StdNormal, StdUniform
+using MeasureBase: mbind, mkernel, bindkernel, boundmeasure
+using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, localmeasure
+
+@testset "bind" begin
+    stblrng() = StableRNG(789990641)
+
+    f_β(σ) = pushfwd(Mul(σ + 0.5), StdNormal())
+    α = StdExponential()
+
+    @testset "monadic bind" begin
+        μ = mbind(f_β, α)
+        @test μ isa MeasureBase.Bind
+        @test boundmeasure(μ) === α
+        @test bindkernel(μ) isa MeasureBase.MKernel
+        @test mbind(bindkernel(μ), α) == μ
+        @test mbind(f_β)(α) == μ
+
+        a = rand(stblrng(), Float64, α)
+        b = rand(copy(stblrng()), Float64, μ)  # not comparable directly, just smoke:
+        @test rand(stblrng(), Float64, μ) isa Real
+
+        @test MeasureBase.insupport(μ, 0.4) isa MeasureBase.NoFastInsupport
+        @test MeasureBase.getdof(μ) isa MeasureBase.NoDOF
+        @test_throws ArgumentError basemeasure(μ)
+        @test_throws ArgumentError MeasureBase.rootmeasure(μ)
+    end
+
+    @testset "mbind with tuple and Pair" begin
+        for f_c in (tuple, Pair)
+            μ = mbind(f_β, α, f_c)
+            ab = rand(stblrng(), Float64, μ)
+            a, b = f_c === tuple ? ab : (ab.first, ab.second)
+            @test logdensityof(μ, ab) ≈ logdensityof(α, a) + logdensityof(f_β(a), b)
+
+            tpm = transportmeasure(μ, ab)
+            @test logdensityof(tpm, ab) ≈ logdensityof(μ, ab)
+            @test localmeasure(μ, ab) == tpm
+        end
+    end
+
+    @testset "mbind with vcat" begin
+        αv = StdExponential()^1
+        f_βv(a) = pushfwd(Mul(a[1] + 0.5), StdNormal())^2
+        μ = mbind(f_βv, αv, vcat)
+
+        xy = rand(stblrng(), Float64, μ)
+        @test xy isa AbstractVector{<:Real} && length(xy) == 3
+        a, b = xy[1:1], xy[2:3]
+        @test logdensityof(μ, xy) ≈ logdensityof(αv, a) + logdensityof(f_βv(a), b)
+
+        # Transport to and from a standard measure, dof of μ is not fast-computable:
+        y = transport_to(StdUniform()^3, μ)(xy)
+        @test y isa AbstractVector{<:Real} && length(y) == 3
+        @test all(u -> 0 <= u <= 1, y)
+        xy_reco = transport_to(μ, StdUniform()^3)(y)
+        @test xy_reco ≈ xy
+    end
+
+    @testset "mbind with merge" begin
+        αnt = productmeasure((position = StdNormal(),))
+        f_βnt(a) = productmeasure((
+            noise = pushfwd(Mul(abs(a.position) + 0.5), StdExponential()),
+        ))
+        μ = mbind(f_βnt, αnt, merge)
+
+        x = rand(stblrng(), Float64, μ)
+        @test x isa NamedTuple{(:position, :noise)}
+        @test logdensityof(μ, x) ≈
+              logdensityof(αnt, (position = x.position,)) +
+              logdensityof(f_βnt(x), (noise = x.noise,))
+
+        y = transport_to(StdUniform()^2, μ)(x)
+        @test y isa AbstractVector{<:Real} && length(y) == 2
+        x_reco = transport_to(μ, StdUniform()^2)(y)
+        @test x_reco.position ≈ x.position && x_reco.noise ≈ x.noise
+    end
+
+    @testset "mbind with Dirac" begin
+        @test mbind(f_β, MeasureBase.Dirac(1.5)) == asmeasure(f_β(1.5))
+        μ = mbind(f_β, MeasureBase.Dirac(1.5), tuple)
+        ab = rand(stblrng(), Float64, μ)
+        @test ab[1] == 1.5
+    end
+end
diff --git a/test/combinators/combined.jl b/test/combinators/combined.jl
new file mode 100644
index 00000000..9a79f225
--- /dev/null
+++ b/test/combinators/combined.jl
@@ -0,0 +1,52 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using Random
+using StableRNGs: StableRNG
+using OneTwoMany: firstarg, secondarg
+
+using MeasureBase
+using MeasureBase: StdExponential, StdLogistic, StdNormal, StdUniform
+using MeasureBase: mcombine, productmeasure, transport_to
+
+@testset "mcombine" begin
+    stblrng() = StableRNG(789990641)
+
+    α = StdExponential()
+    β = StdLogistic()
+
+    @testset "combination shortcuts" begin
+        @test mcombine(firstarg, α, β) === α
+        @test mcombine(secondarg, α, β) === β
+        @test mcombine(tuple, α, β) == productmeasure((α, β))
+        @test mcombine(vcat, StdNormal()^2, StdNormal()^1) == StdNormal()^3
+        @test mcombine(vcat, productmeasure([α, α]), productmeasure([α])) ==
+              productmeasure([α, α, α])
+        @test mcombine(
+            merge,
+            productmeasure((a = α,)),
+            productmeasure((b = β,)),
+        ) == productmeasure((a = α, b = β))
+        @test mcombine(tuple, MeasureBase.Dirac(1), MeasureBase.Dirac(2)) ==
+              MeasureBase.Dirac((1, 2))
+    end
+
+    @testset "CombinedMeasure" begin
+        μ = mcombine(Pair, α, β)
+        @test μ isa MeasureBase.CombinedMeasure
+
+        ab = rand(stblrng(), Float64, μ)
+        @test ab isa Pair
+        @test logdensityof(μ, ab) ≈ logdensityof(α, ab.first) + logdensityof(β, ab.second)
+
+        @test MeasureBase.getdof(μ) == 2
+        @test MeasureBase.fast_dof(μ) == 2
+        @test MeasureBase.insupport(μ, ab) isa MeasureBase.NoFastInsupport
+
+        y = transport_to(StdUniform()^2, μ)(ab)
+        @test y isa AbstractVector{<:Real} && length(y) == 2
+        ab_reco = transport_to(μ, StdUniform()^2)(y)
+        @test ab_reco.first ≈ ab.first && ab_reco.second ≈ ab.second
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 7183c58f..d24f58e2 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -26,6 +26,8 @@ include("combinators/weighted.jl")
 include("combinators/transformedmeasure.jl")
 include("combinators/reshape.jl")
 include("combinators/implicitlymapped.jl")
+include("combinators/combined.jl")
+include("combinators/bind.jl")
 
 include("distributions/test_distributions.jl")
 

From 30973a7267a06d65b8d1914e13d34eb850430576 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 15:48:08 +0200
Subject: [PATCH 044/122] Extend proxy forwarding, pushforward and mass
 interface

Forward `unsafe_logdensityof`, `rootmeasure`, `insupport`, `getdof`,
`fast_dof`, `localmeasure` and `transportmeasure` through `@useproxy`.
Pushforward measures now return `NoFastInsupport` from `insupport`
(checking via the origin would require a potentially costly
transformation), support `fast_dof` and gain curried `pushfwd(f)`/
`pullbck(f)` forms. `massof(::AbstractMeasure)` returns `UnknownMass()`
(instead of an invalid constructor call), unknown-mass powers
disambiguate integer and rational exponents, and the root-measure
based `massof` implementation hook is now named `_default_massof_impl`
and uses interval endpoints.

Ports parts of the original measure-algebra commits 9517fb1, 8aec2f2
and related STASH commits.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/combinators/transformedmeasure.jl  |  9 ++++++++-
 src/mass-interface.jl                  |  9 ++++++---
 src/primitives/lebesgue.jl             |  6 +++---
 src/proxies.jl                         | 11 ++++++++++-
 test/combinators/transformedmeasure.jl |  2 +-
 test/distributions/test_conversions.jl |  4 ++--
 6 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index c9db7a6b..56112049 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -170,7 +170,9 @@ for func in [:logdensityof, :logdensity_def]
     end
 end
 
-insupport(m::PushforwardMeasure, x) = insupport(transport_origin(m), to_origin(m, x))
+# Checking insupport via the origin would require a potentially costly
+# transformation of x:
+insupport(m::PushforwardMeasure, x) = NoFastInsupport{typeof(m)}()
 
 function testvalue(::Type{T}, ν::PushforwardMeasure) where {T}
     ν.f(testvalue(T, parent(ν)))
@@ -193,6 +195,9 @@ _pushfwd_dof(::Type{MU}, ::Type{<:Tuple{Any,Real}}, dof) where {MU} = dof
 @inline getdof(ν::MU) where {MU<:PushforwardMeasure} = getdof(ν.origin)
 @inline getdof(m::_NonBijectivePusfwdMeasure) = MeasureBase.NoDOF{typeof(m)}()
 
+@inline fast_dof(ν::PushforwardMeasure) = fast_dof(ν.origin)
+@inline fast_dof(m::_NonBijectivePusfwdMeasure) = MeasureBase.NoDOF{typeof(m)}()
+
 # Bypass `checked_arg`, would require potentially costly transformation:
 @inline checked_arg(::PushforwardMeasure, x) = x
 
@@ -222,6 +227,7 @@ To manually specify an inverse, call
 function pushfwd end
 export pushfwd
 
+@inline pushfwd(f) = Base.Fix1(pushfwd, f)
 @inline pushfwd(f, μ) = _pushfwd_impl(f, μ, AdaptRootMeasure())
 @inline pushfwd(f, μ, style::AdaptRootMeasure) = _pushfwd_impl(f, μ, style)
 @inline pushfwd(f, μ, style::PushfwdRootMeasure) = _pushfwd_impl(f, μ, style)
@@ -263,6 +269,7 @@ To manually specify an inverse, call
 function pullbck end
 export pullbck
 
+@inline pullbck(f) = Base.Fix1(pullbck, f)
 @inline pullbck(f, μ) = _pullback_impl(f, μ, AdaptRootMeasure())
 @inline pullbck(f, μ, style::AdaptRootMeasure) = _pullback_impl(f, μ, style)
 @inline pullbck(f, μ, style::PushfwdRootMeasure) = _pullback_impl(f, μ, style)
diff --git a/src/mass-interface.jl b/src/mass-interface.jl
index 7b0518f9..59a5db88 100644
--- a/src/mass-interface.jl
+++ b/src/mass-interface.jl
@@ -22,7 +22,10 @@ for T in (:UnknownFiniteMass, :UnknownMass)
     @eval begin
         Base.:+(::$T, ::$T) = $T()
         Base.:*(::$T, ::$T) = $T()
-        Base.:^(::$T, k::Number) = isfinite(k) ? $T() : UnknownMass()
+        Base.:^(::$T, k::Real) = isfinite(k) ? $T() : UnknownMass()
+        # Disambiguation:
+        Base.:^(::$T, k::Integer) = isfinite(k) ? $T() : UnknownMass()
+        Base.:^(::$T, k::Rational) = isfinite(k) ? $T() : UnknownMass()
     end
 end
 
@@ -65,7 +68,7 @@ finite, or we may know nothing at all about it. For these cases, it will return
 `UnknownFiniteMass` or `UnknownMass`, respectively. When no `massof` method
 exists, it defaults to `UnknownMass`.
 """
-massof(m::AbstractMeasure) = UnknownMass(m)
+massof(::AbstractMeasure) = UnknownMass()
 
 struct NormalizedMeasure{P,M} <: AbstractMeasure
     parent::P
@@ -105,7 +108,7 @@ isnormalized(x, p::Real = 2) = isone(norm(x, p))
 isone(::AbstractUnknownMass) = false
 
 function massof(m, s)
-    _massof(m, s, rootmeasure(m))
+    _default_massof_impl(m, s, rootmeasure(m))
 end
 
 """
diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl
index 3d92a2ed..4b8bf7ab 100644
--- a/src/primitives/lebesgue.jl
+++ b/src/primitives/lebesgue.jl
@@ -26,12 +26,12 @@ end
 
 massof(::LebesgueBase) = static(Inf)
 
-function _massof(m, s::Interval, ::LebesgueBase)
+function _default_massof_impl(m, s::AbstractInterval, ::LebesgueBase)
     mass = massof(m)
     nu = mass * StdUniform()
     f = transport_to(nu, m)
-    a = f(minimum(s))
-    b = f(maximum(s))
+    a = f(leftendpoint(s))
+    b = f(rightendpoint(s))
     return mass * abs(b - a)
 end
 
diff --git a/src/proxies.jl b/src/proxies.jl
index 95aed270..109b9973 100644
--- a/src/proxies.jl
+++ b/src/proxies.jl
@@ -15,15 +15,24 @@ macro useproxy(M)
     M = esc(M)
     quote
         @inline $MeasureBase.logdensity_def(μ::$M, x) = logdensity_def(proxy(μ), x)
+        @inline $MeasureBase.unsafe_logdensityof(μ::$M, x) = unsafe_logdensityof(proxy(μ), x)
 
         @inline $MeasureBase.basemeasure(μ::$M) = basemeasure(proxy(μ))
-
         @inline $MeasureBase.basemeasure_depth(μ::$M) = basemeasure_depth(proxy(μ))
+        @inline $MeasureBase.rootmeasure(μ::$M) = rootmeasure(proxy(μ))
+
+        @inline $MeasureBase.insupport(μ::$M) = insupport(proxy(μ))
+
+        @inline $MeasureBase.getdof(μ::$M) = getdof(proxy(μ))
+        @inline $MeasureBase.fast_dof(μ::$M) = fast_dof(proxy(μ))
 
         @inline $MeasureBase.transport_origin(μ::$M) = transport_origin(proxy(μ))
         @inline $MeasureBase.to_origin(μ::$M, y) = to_origin(proxy(μ), y)
         @inline $MeasureBase.from_origin(μ::$M, x) = from_origin(proxy(μ), x)
 
+        @inline $MeasureBase.localmeasure(μ::$M, x) = localmeasure(proxy(μ), x)
+        @inline $MeasureBase.transportmeasure(μ::$M, x) = transportmeasure(proxy(μ), x)
+
         @inline $MeasureBase.massof(μ::$M) = massof(proxy(μ))
         @inline $MeasureBase.massof(μ::$M, s) = massof(proxy(μ), s)
 
diff --git a/test/combinators/transformedmeasure.jl b/test/combinators/transformedmeasure.jl
index 28ddbb50..497f79c2 100644
--- a/test/combinators/transformedmeasure.jl
+++ b/test/combinators/transformedmeasure.jl
@@ -162,7 +162,7 @@ using ChangesOfVariables
 
     # Test rand
     @test rand(ν) isa Real
-    @test insupport(ν, rand(ν))
+    @test insupport(ν, rand(ν)) != false
 
     # Test pullback
     pb = pullbck(f, ν)
diff --git a/test/distributions/test_conversions.jl b/test/distributions/test_conversions.jl
index 0f9c4d80..a8244b4d 100644
--- a/test/distributions/test_conversions.jl
+++ b/test/distributions/test_conversions.jl
@@ -30,13 +30,13 @@ using MeasureBase: logdensityof, massof, insupport
             for x in (rand(stblrng(), d) for _ in 1:10)
                 @test logdensityof(m, x) ≈ logpdf(d, x)
                 @test logpdf(d2, x) ≈ logpdf(d, x)
-                @test insupport(m, x)
+                @test insupport(m, x) != false
             end
 
             x = rand(stblrng(), Float64, m)
             # Tuple-marginal product measures have tuple variates:
             x isa Tuple ? (@test length(x) == length(d)) : (@test size(x) == size(d))
-            @test insupport(m, x)
+            @test insupport(m, x) != false
         end
     end
 

From 27863de892ec77982ac5c8eea0cb8b4b18b13037 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 17:43:00 +0200
Subject: [PATCH 045/122] Use PushFwdStyle names directly, keep VolCorr names
 as compat aliases only

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/combinators/transformedmeasure.jl | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 56112049..970b862d 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -8,6 +8,7 @@ pushforward. Either [`AdaptRootMeasure()`](@ref) or
 abstract type PushFwdStyle end
 export PushFwdStyle
 
+# Backward compatibility with user code, do not use in MeasureBase itself:
 const TransformVolCorr = PushFwdStyle
 
 """
@@ -22,9 +23,10 @@ Density calculations for pushforward measures constructed with
 transform (typically via the log-abs-det-Jacobian of the transform) into
 account.
 """
-struct AdaptRootMeasure <: TransformVolCorr end
+struct AdaptRootMeasure <: PushFwdStyle end
 export AdaptRootMeasure
 
+# Backward compatibility with user code, do not use in MeasureBase itself:
 const WithVolCorr = AdaptRootMeasure
 
 """
@@ -37,9 +39,10 @@ Density calculations for pushforward measures constructed with
 `PushfwdRootMeasure()` will ignore the volume element of the variate
 transform.
 """
-struct PushfwdRootMeasure <: TransformVolCorr end
+struct PushfwdRootMeasure <: PushFwdStyle end
 export PushfwdRootMeasure
 
+# Backward compatibility with user code, do not use in MeasureBase itself:
 const NoVolCorr = PushfwdRootMeasure
 
 abstract type AbstractTransformedMeasure <: AbstractMeasure end

From 451c95d019f38b0000557c840eeb93b866d899f6 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 20:02:13 +0200
Subject: [PATCH 046/122] Fix logweight scaling in powers of weighted measures

The total weight was computed from the length of the new power measure,
which fails when the base collapses to a non-power measure (e.g. for
Dirac bases). Compute it from the exponent instead.

Created by generative AI.
---
 src/combinators/smart-constructors.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index cb15f6ad..a7fb9a80 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -41,7 +41,7 @@ end
 
 @inline function _generic_powermeasure_stage2(μ::WeightedMeasure, exponent::Tuple)
     ν = μ.base^exponent
-    k = maybestatic_length(ν) * μ.logweight
+    k = size2length(axes2size(exponent)) * μ.logweight
     return weightedmeasure(k, ν)
 end
 

From da2ce09680c93d0903241bfb72e93f545f4653f4 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 20:02:44 +0200
Subject: [PATCH 047/122] Rework superpose into an optimizing pairwise algebra

superpose now folds varargs pairwise: equal measures combine into
weighted measures, weighted measures with equal bases add their weights,
and superpositions merge their components. Merging no longer mutates
existing superposition components (add_measures used push!). Collections
of a single singleton measure type collapse to weighted measures.

Also exports superpose, matching the other smart constructors.

Created by generative AI.
---
 src/collection_utils.jl               |  7 +++
 src/combinators/smart-constructors.jl | 69 ++++++++++++++++++++-------
 2 files changed, 60 insertions(+), 16 deletions(-)

diff --git a/src/collection_utils.jl b/src/collection_utils.jl
index d15131d0..be5cc934 100644
--- a/src/collection_utils.jl
+++ b/src/collection_utils.jl
@@ -94,3 +94,10 @@ _flatten_to_rv(VV::VectorOfVectors{<:Real}) = flatview(VV)
 _flatten_to_rv(::Tuple{}) = []
 _flatten_to_rv(tpl::Tuple{Vararg{AbstractVector}}) = vcat(tpl...)
 _flatten_to_rv(tpl::Tuple{Vararg{StaticVector}}) = vcat(tpl...)
+
+
+# Non-mutating concatenation of measure collections:
+_cat_measures(a::Tuple, b::Tuple) = (a..., b...)
+_cat_measures(a::AbstractVector, b::Tuple) = vcat(a, [b...])
+_cat_measures(a::Tuple, b::AbstractVector) = vcat([a...], b)
+_cat_measures(a::AbstractVector, b::AbstractVector) = vcat(a, b)
diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index a7fb9a80..2b52fede 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -134,35 +134,72 @@ restrict(f, b) = RestrictedMeasure(f, b)
 ###############################################################################
 # SuperpositionMeasure
 
-superpose(a::AbstractArray) = SuperpositionMeasure(a)
+"""
+    superpose(μs...)
+    superpose(μs)
 
-superpose(t::Tuple) = SuperpositionMeasure(t)
-superpose(nt::NamedTuple) = SuperpositionMeasure(nt)
+Constructs a superposition of measures, given either as separate arguments or
+as a collection (array, tuple or named tuple) of measures.
+
+The vararg form simplifies algebraically: equal measures combine into weighted
+measures (`superpose(μ, μ) == weightedmeasure(log(2), μ)`), weighted measures
+with equal bases add their weights, and superpositions merge their components.
+Collections are wrapped as-is, apart from cost-free structural simplifications.
+"""
+function superpose end
+export superpose
+
+superpose(μ::AbstractMeasure) = μ
+
+function superpose(μ::AbstractMeasure, ν::AbstractMeasure, more::AbstractMeasure...)
+    superpose(_superpose_two(μ, ν), more...)
+end
 
-function superpose(μ::T, ν::T) where {T<:AbstractMeasure}
-    if μ == ν
-        return weightedmeasure(static(float(logtwo)), μ)
+function superpose(a::AbstractArray{T}) where {T}
+    if Base.issingletontype(T)
+        weightedmeasure(log(length(a)), asmeasure(instance(T)))
     else
-        return superpose((μ, ν))
+        SuperpositionMeasure(a)
     end
 end
 
-function superpose(μ::AbstractMeasure, μs...)
-    if all(==(μ), μs)
-        return weightedmeasure(log(length(μs) + 1), μ)
+superpose(a::FillArrays.Fill) = weightedmeasure(log(length(a)), asmeasure(_fill_value(a)))
+
+superpose(t::Tuple) = SuperpositionMeasure(t)
+superpose(nt::NamedTuple) = SuperpositionMeasure(nt)
+
+function _superpose_two(μ::AbstractMeasure, ν::AbstractMeasure)
+    μ == ν ? weightedmeasure(static(float(logtwo)), μ) : SuperpositionMeasure((μ, ν))
+end
+
+function _superpose_two(μ::WeightedMeasure, ν::WeightedMeasure)
+    if μ.base == ν.base
+        weightedmeasure(logaddexp(asnonstatic(μ.logweight), asnonstatic(ν.logweight)), μ.base)
     else
-        return superpose((μ, μs...))
+        SuperpositionMeasure((μ, ν))
     end
 end
 
-add_measures(μs::AbstractVector, νs) = push!(μs, νs...)
-add_measures(μs::Tuple, νs) = (μs..., νs...)
+function _superpose_two(μ::WeightedMeasure, ν::AbstractMeasure)
+    μ.base == ν ? weightedmeasure(log1pexp(asnonstatic(μ.logweight)), μ.base) :
+    SuperpositionMeasure((μ, ν))
+end
 
-function superpose(μ::SuperpositionMeasure, μs...)
-    SuperpositionMeasure(add_measures(μ.components, μs))
+function _superpose_two(μ::AbstractMeasure, ν::WeightedMeasure)
+    μ == ν.base ? weightedmeasure(log1pexp(asnonstatic(ν.logweight)), ν.base) :
+    SuperpositionMeasure((μ, ν))
 end
 
-superpose(μ::SuperpositionMeasure) = μ
+_superpose_two(μ::SuperpositionMeasure, ν::SuperpositionMeasure) =
+    SuperpositionMeasure(_cat_measures(μ.components, ν.components))
+_superpose_two(μ::SuperpositionMeasure, ν::AbstractMeasure) =
+    SuperpositionMeasure(_cat_measures(μ.components, (ν,)))
+_superpose_two(μ::AbstractMeasure, ν::SuperpositionMeasure) =
+    SuperpositionMeasure(_cat_measures((μ,), ν.components))
+_superpose_two(μ::SuperpositionMeasure, ν::WeightedMeasure) =
+    SuperpositionMeasure(_cat_measures(μ.components, (ν,)))
+_superpose_two(μ::WeightedMeasure, ν::SuperpositionMeasure) =
+    SuperpositionMeasure(_cat_measures((μ,), ν.components))
 
 ###############################################################################
 # WeightedMeasure

From 4263d57beedd75502fb779574f862f711fa655ba Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 20:02:53 +0200
Subject: [PATCH 048/122] Add value-based equality and isapprox for Dirac

Dirac measures with equal points are equal measures; the default
object identity failed for array-valued points.

Created by generative AI.
---
 src/primitives/dirac.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 01297486..f5e5931a 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -9,6 +9,9 @@ function Pretty.tile(d::Dirac)
     Pretty.literal("Dirac(") * Pretty.tile(d.x) * Pretty.literal(")")
 end
 
+Base.:(==)(a::Dirac, b::Dirac) = a.x == b.x
+Base.isapprox(a::Dirac, b::Dirac; kwargs...) = isapprox(a.x, b.x; kwargs...)
+
 gentype(μ::Dirac{X}) where {X} = X
 
 function (μ::Dirac{X})(s) where {X}

From d7b63b25e3217b4e5aafbe20b1e49cfac60b1f49 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 20:03:26 +0200
Subject: [PATCH 049/122] Extend product measure simplifications

Products of weighted measures (tuples, named tuples and arrays) now pull
the total weight out, following the canonical measure nesting. Empty
products are explicit unit measures (Dirac of the empty variate). Arrays
of measures are no longer copied via broadcast asmeasure, and arrays of
a single singleton measure type collapse to power measures without
requiring a non-empty array. The static-weight optimization now stays
static when the array length is static.

Created by generative AI.
---
 src/combinators/smart-constructors.jl | 46 +++++++++++++++++----------
 1 file changed, 30 insertions(+), 16 deletions(-)

diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index 2b52fede..b3a00496 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -69,29 +69,50 @@ export productmeasure
 @inline _generic_productmeasure_impl(mar::FillArrays.Fill) =
     powermeasure(_fill_value(mar), _fill_axes(mar))
 
+# Empty products are unit measures:
+@inline _generic_productmeasure_impl(::Tuple{}) = Dirac(())
+@inline _generic_productmeasure_impl(::NamedTuple{()}) = Dirac(NamedTuple())
+
 @inline _generic_productmeasure_impl(mar::Tuple{Vararg{AbstractMeasure}}) =
     ProductMeasure(mar)
-_generic_productmeasure_impl(mar::Tuple{Vararg{Dirac}}) = Dirac(map(m -> m.x, mar))
+_generic_productmeasure_impl(mar::Tuple{Dirac,Vararg{Dirac}}) = Dirac(map(m -> m.x, mar))
+_generic_productmeasure_impl(mar::Tuple{WeightedMeasure,Vararg{WeightedMeasure}}) =
+    weightedmeasure(sum(map(_logweight, mar)), productmeasure(map(m -> m.base, mar)))
 _generic_productmeasure_impl(mar::Tuple) = productmeasure(map(asmeasure, mar))
 
 @inline _generic_productmeasure_impl(
     mar::NamedTuple{names,<:Tuple{Vararg{AbstractMeasure}}},
 ) where {names} = ProductMeasure(mar)
-_generic_productmeasure_impl(mar::NamedTuple{names,<:Tuple{Vararg{Dirac}}}) where {names} =
-    Dirac(map(m -> m.x, mar))
+_generic_productmeasure_impl(
+    mar::NamedTuple{names,<:Tuple{Dirac,Vararg{Dirac}}},
+) where {names} = Dirac(map(m -> m.x, mar))
+_generic_productmeasure_impl(
+    mar::NamedTuple{names,<:Tuple{WeightedMeasure,Vararg{WeightedMeasure}}},
+) where {names} =
+    weightedmeasure(sum(map(_logweight, values(mar))), productmeasure(map(m -> m.base, mar)))
 _generic_productmeasure_impl(mar::NamedTuple) = productmeasure(map(asmeasure, mar))
 
-@inline _generic_productmeasure_impl(mar::AbstractArray{<:AbstractProductMeasure}) =
-    ProductMeasure(mar)
-
 _generic_productmeasure_impl(mar::AbstractArray{<:Dirac}) = Dirac((m -> m.x).(mar))
 
-# TODO: We should be able to further optimize this
+_generic_productmeasure_impl(mar::AbstractArray{<:WeightedMeasure}) =
+    weightedmeasure(sum(_logweight, mar), productmeasure((m -> m.base).(mar)))
+
+@inline function _generic_productmeasure_impl(
+    mar::AbstractArray{<:WeightedMeasure{StaticFloat64{W},M}},
+) where {W,M}
+    return weightedmeasure(
+        static(W) * maybestatic_length(mar),
+        productmeasure((m -> m.base).(mar)),
+    )
+end
+
 function _generic_productmeasure_impl(mar::AbstractArray{T}) where {T}
     if Base.issingletontype(T)
-        first(mar)^size(mar)
+        powermeasure(instance(T), axes(mar))
+    elseif T <: AbstractMeasure
+        ProductMeasure(mar)
     else
-        ProductMeasure(asmeasure.(mar))
+        ProductMeasure(map(asmeasure, mar))
     end
 end
 
@@ -103,13 +124,6 @@ end
 
 @inline _generic_productmeasure_impl(mar::Base.Generator) = ProductMeasure(mar)
 
-# TODO: Make this static when its length is static
-@inline function _generic_productmeasure_impl(
-    mar::AbstractArray{<:WeightedMeasure{StaticFloat64{W},M}},
-) where {W,M}
-    return weightedmeasure(W * length(mar), productmeasure(map(basemeasure, mar)))
-end
-
 # ToDo: Remove or at least refactor this (ProductMeasure shouldn't take a kernel as its argument).
 
 productmeasure(f, param_maps, pars) = productmeasure(kernel(f, param_maps), pars)

From 32e5913ef96ec373be31c6c96abf64ad2dc8407b Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 20:03:50 +0200
Subject: [PATCH 050/122] Simplify pushforwards of Dirac and weighted measures

Pushforwards of point masses collapse to point masses and pushforwards
commute with weighting. The identity shortcut moves from _pushfwd_impl
to pushfwd itself and the two style-specific pushfwd/pullbck methods
merge into single PushFwdStyle methods, keeping dispatch unambiguous.

Created by generative AI.
---
 src/combinators/smart-constructors.jl  | 13 +++++++++++++
 src/combinators/transformedmeasure.jl  | 13 +++++++------
 test/combinators/transformedmeasure.jl |  7 ++++---
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index b3a00496..90b646c8 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -139,6 +139,19 @@ function productmeasure(f::Returns{W}, ::typeof(identity), pars) where {W<:Weigh
     weightedmeasure(length(pars) * ℓ, newbase)
 end
 
+###############################################################################
+# PushforwardMeasure
+
+# The pushforward of a point mass is a point mass. Note that no density
+# volume correction applies, Dirac measures are density-defined relative
+# to counting measure:
+_pushfwd_impl(f, μ::Dirac, ::PushFwdStyle) = Dirac(f(μ.x))
+
+# Pushforward and weighting commute:
+function _pushfwd_impl(f, μ::WeightedMeasure, style::PushFwdStyle)
+    weightedmeasure(μ.logweight, _pushfwd_impl(f, μ.base, style))
+end
+
 ###############################################################################
 # RestrictedMeasure
 export restrict
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 970b862d..0017c86f 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -232,8 +232,10 @@ export pushfwd
 
 @inline pushfwd(f) = Base.Fix1(pushfwd, f)
 @inline pushfwd(f, μ) = _pushfwd_impl(f, μ, AdaptRootMeasure())
-@inline pushfwd(f, μ, style::AdaptRootMeasure) = _pushfwd_impl(f, μ, style)
-@inline pushfwd(f, μ, style::PushfwdRootMeasure) = _pushfwd_impl(f, μ, style)
+@inline pushfwd(f, μ, style::PushFwdStyle) = _pushfwd_impl(f, μ, style)
+
+@inline pushfwd(::typeof(identity), μ) = μ
+@inline pushfwd(::typeof(identity), μ, ::PushFwdStyle) = μ
 
 _pushfwd_impl(f, μ, style) = PushforwardMeasure(f, inverse(f), μ, style)
 
@@ -248,8 +250,8 @@ function _pushfwd_impl(
     PushforwardMeasure(new_f, new_f_inv, orig_μ, style)
 end
 
-_pushfwd_impl(::typeof(identity), μ, ::AdaptRootMeasure) = μ
-_pushfwd_impl(::typeof(identity), μ, ::PushfwdRootMeasure) = μ
+# Simplifications for Dirac and WeightedMeasure origins are defined in
+# smart-constructors.jl.
 
 ###############################################################################
 # pullback
@@ -274,8 +276,7 @@ export pullbck
 
 @inline pullbck(f) = Base.Fix1(pullbck, f)
 @inline pullbck(f, μ) = _pullback_impl(f, μ, AdaptRootMeasure())
-@inline pullbck(f, μ, style::AdaptRootMeasure) = _pullback_impl(f, μ, style)
-@inline pullbck(f, μ, style::PushfwdRootMeasure) = _pullback_impl(f, μ, style)
+@inline pullbck(f, μ, style::PushFwdStyle) = _pullback_impl(f, μ, style)
 
 function _pullback_impl(f, μ, style = AdaptRootMeasure())
     pushfwd(inverse(f), μ, style)
diff --git a/test/combinators/transformedmeasure.jl b/test/combinators/transformedmeasure.jl
index 497f79c2..6261c762 100644
--- a/test/combinators/transformedmeasure.jl
+++ b/test/combinators/transformedmeasure.jl
@@ -152,9 +152,10 @@ using ChangesOfVariables
     @test rootmeasure(ν) === rootmeasure(μ)  # AdaptRootMeasure
     @test rootmeasure(ν_no_corr) isa PushforwardMeasure  # PushfwdRootMeasure
 
-    # Test basemeasure
-    @test basemeasure(ν) isa PushforwardMeasure
-    @test basemeasure(ν).style isa PushfwdRootMeasure
+    # Test basemeasure. The base measure of μ is a weighted Lebesgue measure,
+    # so the weight gets pulled out of the pushforward:
+    @test basemeasure(ν) isa WeightedMeasure{<:Any,<:PushforwardMeasure}
+    @test basemeasure(ν).base.style isa PushfwdRootMeasure
 
     # Test massof
     # TODO: mass interface is very incomplete

From 7ec90978e4b37eee01a6ce00093d5fcbcf0eb04e Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 20:04:07 +0200
Subject: [PATCH 051/122] Fuse nested measure restrictions and add curried
 restrict

Created by generative AI.
---
 src/combinators/smart-constructors.jl | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index 90b646c8..8292fa1c 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -156,7 +156,13 @@ end
 # RestrictedMeasure
 export restrict
 
-restrict(f, b) = RestrictedMeasure(f, b)
+@inline restrict(f) = Base.Fix1(restrict, f)
+
+restrict(f, μ) = RestrictedMeasure(f, asmeasure(μ))
+
+# Nested restrictions fuse into a single predicate:
+restrict(f, μ::RestrictedMeasure) =
+    RestrictedMeasure(x -> μ.predicate(x) && f(x), μ.base)
 
 ###############################################################################
 # SuperpositionMeasure

From af701872bf4e4c123d1ed9f4670766ef9e5df07d Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 20:04:29 +0200
Subject: [PATCH 052/122] Document and export half and weightedmeasure

Matches the other smart constructors, which are documented and exported.

Created by generative AI.
---
 src/combinators/smart-constructors.jl | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index 8292fa1c..db355184 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -7,7 +7,14 @@
 ###############################################################################
 # Half
 
+"""
+    half(μ::AbstractMeasure)
+
+Constructs the half-measure of a measure `μ` that is symmetric around zero:
+`μ` folded onto the non-negative half-line.
+"""
 half(μ::AbstractMeasure) = Half(μ)
+export half
 
 ###############################################################################
 # PowerMeaure
@@ -237,6 +244,16 @@ _superpose_two(μ::WeightedMeasure, ν::SuperpositionMeasure) =
 ###############################################################################
 # WeightedMeasure
 
+"""
+    weightedmeasure(logweight::Real, μ)
+
+Constructs a measure that behaves like the measure `μ`, but with its density
+scaled by `exp(logweight)`. Weights of nested weighted measures combine
+additively.
+"""
+function weightedmeasure end
+export weightedmeasure
+
 function weightedmeasure(ℓ::R, b::M) where {R,M}
     WeightedMeasure{R,M}(ℓ, b)
 end

From 0ce4beb8056209889655232510fb5a0f65e0cc1d Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 20:04:44 +0200
Subject: [PATCH 053/122] Add smart constructor tests

Covers the power/product/superpose/pushfwd/restrict simplifications.
Also wires the previously orphaned superpose tests into the test suite,
updated to current basemeasure collapse behavior.

Created by generative AI.
---
 test/combinators/smart_constructors.jl | 139 +++++++++++++++++++++++++
 test/combinators/superpose.jl          |   5 +-
 test/runtests.jl                       |   2 +
 3 files changed, 143 insertions(+), 3 deletions(-)
 create mode 100644 test/combinators/smart_constructors.jl

diff --git a/test/combinators/smart_constructors.jl b/test/combinators/smart_constructors.jl
new file mode 100644
index 00000000..ac6da161
--- /dev/null
+++ b/test/combinators/smart_constructors.jl
@@ -0,0 +1,139 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using MeasureBase
+using MeasureBase:
+    weightedmeasure, superpose, productmeasure, powermeasure, pushfwd, pullbck, restrict
+using MeasureBase:
+    WeightedMeasure,
+    SuperpositionMeasure,
+    ProductMeasure,
+    PowerMeasure,
+    PushforwardMeasure,
+    RestrictedMeasure,
+    Dirac,
+    StdNormal,
+    StdUniform,
+    StdExponential,
+    PushfwdRootMeasure,
+    AdaptRootMeasure
+using FillArrays: Fill
+using Static: static
+
+@testset "smart constructors" begin
+    @testset "powermeasure" begin
+        @test powermeasure(StdNormal(), ()) === StdNormal()
+
+        @test powermeasure(Dirac(4.2), (3,)) == Dirac(Fill(4.2, 3))
+
+        wpw = weightedmeasure(0.3, StdNormal())^(2, 3)
+        @test wpw isa WeightedMeasure
+        @test wpw.logweight ≈ 6 * 0.3
+        @test wpw.base == StdNormal()^(2, 3)
+
+        # Weight pull-out must not depend on the collapsed base type:
+        wd = weightedmeasure(0.3, Dirac(1.5))^3
+        @test wd isa WeightedMeasure
+        @test wd.logweight ≈ 3 * 0.3
+        @test wd.base == Dirac(Fill(1.5, 3))
+
+        ws = weightedmeasure(static(0.5), StdNormal())^static(4)
+        @test ws.logweight ≈ 2.0
+    end
+
+    @testset "productmeasure" begin
+        @test productmeasure(Fill(StdUniform(), 3)) == StdUniform()^3
+
+        @test productmeasure(()) === Dirac(())
+        @test productmeasure(NamedTuple()) === Dirac(NamedTuple())
+
+        @test productmeasure((Dirac(1), Dirac(2))) === Dirac((1, 2))
+        @test productmeasure((a = Dirac(1), b = Dirac(2))) === Dirac((a = 1, b = 2))
+        @test productmeasure([Dirac(1), Dirac(2)]) == Dirac([1, 2])
+
+        pt = productmeasure((2.0 * StdNormal(), 3.0 * StdUniform()))
+        @test pt isa WeightedMeasure
+        @test exp(pt.logweight) ≈ 6
+        @test pt.base == ProductMeasure((StdNormal(), StdUniform()))
+
+        pnt = productmeasure((a = 2.0 * StdNormal(), b = 3.0 * StdUniform()))
+        @test pnt isa WeightedMeasure
+        @test exp(pnt.logweight) ≈ 6
+        @test pnt.base == ProductMeasure((a = StdNormal(), b = StdUniform()))
+
+        pa = productmeasure([2.0 * StdNormal(), 3.0 * StdNormal()])
+        @test pa isa WeightedMeasure
+        @test exp(pa.logweight) ≈ 6
+        @test pa.base == StdNormal()^2
+
+        @test logdensityof(pt, (0.3, 0.5)) ≈ log(6) + logdensityof(StdNormal(), 0.3)
+
+        @test productmeasure([StdNormal(), StdNormal()]) == StdNormal()^2
+        @test productmeasure([StdNormal()^2, StdUniform()^3]) isa ProductMeasure
+    end
+
+    @testset "superpose" begin
+        μ, ν = StdNormal(), StdUniform()
+
+        @test superpose(μ) === μ
+        @test superpose(μ, ν) == SuperpositionMeasure((μ, ν))
+
+        s2 = superpose(μ, μ)
+        @test s2 isa WeightedMeasure && exp(s2.logweight) ≈ 2 && s2.base === μ
+
+        s4 = superpose(μ, μ, μ, μ)
+        @test s4 isa WeightedMeasure && exp(s4.logweight) ≈ 4
+
+        c = superpose(2.0 * μ, 3.0 * μ)
+        @test c isa WeightedMeasure && exp(c.logweight) ≈ 5 && c.base === μ
+        @test exp(superpose(2.0 * μ, μ).logweight) ≈ 3
+        @test exp(superpose(μ, 2.0 * μ).logweight) ≈ 3
+        @test superpose(2.0 * μ, 3.0 * ν) == SuperpositionMeasure((2.0 * μ, 3.0 * ν))
+
+        ss = superpose(superpose((μ, ν)), superpose((ν, StdExponential())))
+        @test ss.components === (μ, ν, ν, StdExponential())
+        @test superpose(superpose((μ, ν)), StdExponential()).components ===
+              (μ, ν, StdExponential())
+        @test superpose(StdExponential(), superpose((μ, ν))).components ===
+              (StdExponential(), μ, ν)
+
+        # Merging must not mutate existing superpositions:
+        sv = superpose(AbstractMeasure[μ, ν])
+        sv2 = superpose(sv, StdExponential())
+        @test length(sv.components) == 2 && length(sv2.components) == 3
+
+        @test superpose(Fill(μ, 4)) == weightedmeasure(log(4), μ)
+        @test superpose([μ, μ, μ]) == weightedmeasure(log(3), μ)
+
+        @test logdensityof(c, 0.3) ≈ log(5) + logdensityof(μ, 0.3)
+    end
+
+    @testset "pushfwd" begin
+        @test pushfwd(identity, StdNormal()) === StdNormal()
+        @test pushfwd(identity, StdNormal(), PushfwdRootMeasure()) === StdNormal()
+
+        @test pushfwd(sqrt, Dirac(4.0)) === Dirac(2.0)
+        @test pushfwd(sqrt, Dirac(4.0), PushfwdRootMeasure()) === Dirac(2.0)
+
+        pw = pushfwd(sqrt, 3.0 * StdExponential())
+        @test pw isa WeightedMeasure && exp(pw.logweight) ≈ 3
+        @test pw.base isa PushforwardMeasure
+
+        pp = pushfwd(exp, pushfwd(sqrt, StdExponential()))
+        @test pp isa PushforwardMeasure && pp.origin === StdExponential()
+
+        @test pullbck(log, Dirac(4.0)) === Dirac(exp(4.0))
+    end
+
+    @testset "restrict" begin
+        r = restrict(x -> x > 0, StdNormal())
+        @test r isa RestrictedMeasure && r.base === StdNormal()
+
+        r2 = restrict(x -> x < 1, r)
+        @test r2 isa RestrictedMeasure && r2.base === StdNormal()
+        @test r2.predicate(0.5) && !r2.predicate(-1.0) && !r2.predicate(2.0)
+
+        @test restrict(x -> x > 0)(StdNormal()) isa RestrictedMeasure
+    end
+end
diff --git a/test/combinators/superpose.jl b/test/combinators/superpose.jl
index ed4c6996..17c677df 100644
--- a/test/combinators/superpose.jl
+++ b/test/combinators/superpose.jl
@@ -14,9 +14,8 @@ using MeasureBase: superpose
 
     μs = SuperpositionMeasure([μ, ν])
     @test μs isa SuperpositionMeasure{<:AbstractVector{<:AbstractMeasure}}
-    @test_throws ErrorException density_def(μs, 0)
-    @test basemeasure(μs).components ==
-          SuperpositionMeasure([CountingBase(), CountingBase()]).components
+    @test density_def(μs, 0) == 1.0
+    @test basemeasure(μs) == weightedmeasure(log(2), CountingBase())
 
     μ2 = μ + μ
     @test μ2 isa WeightedMeasure
diff --git a/test/runtests.jl b/test/runtests.jl
index d24f58e2..6168091a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -22,7 +22,9 @@ include("test_mooncake.jl")
 
 include("measure_operators.jl")
 
+include("combinators/smart_constructors.jl")
 include("combinators/weighted.jl")
+include("combinators/superpose.jl")
 include("combinators/transformedmeasure.jl")
 include("combinators/reshape.jl")
 include("combinators/implicitlymapped.jl")

From 5aa67b8363713bb273caae82fcdd735693fba509 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 20:18:01 +0200
Subject: [PATCH 054/122] Support transport between measures of unknown DOF

Measure transport now falls back to a pivot through a flat vector of
standard uniform variates, using the with-rest transport protocol, when
the DOF of either side is not fast-computable. This enables transport
between hierarchical measures (Bind) and between known-DOF and
unknown-DOF measures, as already promised by the transport_def
docstring.

The mvstd gateway now also verifies that the produced variate length
matches the target standard power measure, instead of silently
returning a variate of the wrong length.

Created by generative AI.
---
 src/combinators/product_transport.jl |  6 +++++-
 src/transport.jl                     | 20 +++++++++++++++++++-
 test/combinators/bind.jl             | 17 +++++++++++++++++
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/src/combinators/product_transport.jl b/src/combinators/product_transport.jl
index 01c1016c..2d78ae2f 100644
--- a/src/combinators/product_transport.jl
+++ b/src/combinators/product_transport.jl
@@ -124,7 +124,11 @@ end
 # Transport to a multivariate standard measure from any measure:
 
 function transport_def(ν::StdPowerMeasure{NU,1}, μ::AbstractMeasure, x) where {NU}
-    transport_to_mvstd(pwr_base(ν), μ, x)
+    y = transport_to_mvstd(pwr_base(ν), μ, x)
+    if maybestatic_length(y) != maybestatic_length(ν)
+        throw(ArgumentError("Length of transport target doesn't match variate DOF during transport"))
+    end
+    return y
 end
 
 function transport_to_mvstd(ν_inner::StdMeasure, μ::AbstractMeasure, x)
diff --git a/src/transport.jl b/src/transport.jl
index 9cea6fed..65c5f386 100644
--- a/src/transport.jl
+++ b/src/transport.jl
@@ -137,11 +137,29 @@ end
     return static(10)
 end
 
-# If both both measures have no origin:
+# If both measures have no origin:
 function _transport_between_origins(ν, ::StaticInteger{0}, ::StaticInteger{0}, μ, x)
+    _transport_between_noorigins(ν, fast_dof(ν), fast_dof(μ), μ, x)
+end
+
+function _transport_between_noorigins(ν, ::IntegerLike, ::IntegerLike, μ, x)
     _transport_with_intermediate(ν, _transport_intermediate(ν, μ), μ, x)
 end
 
+# If the DOF of either side is not known, pivot through a flat vector of
+# standard-measure variates, using the with-rest transport protocol:
+_transport_between_noorigins(ν, ::AbstractNoDOF, ::IntegerLike, μ, x) =
+    _transport_mvstd_pivot(ν, μ, x)
+_transport_between_noorigins(ν, ::IntegerLike, ::AbstractNoDOF, μ, x) =
+    _transport_mvstd_pivot(ν, μ, x)
+_transport_between_noorigins(ν, ::AbstractNoDOF, ::AbstractNoDOF, μ, x) =
+    _transport_mvstd_pivot(ν, μ, x)
+
+function _transport_mvstd_pivot(ν, μ, x)
+    z = transport_to_mvstd(StdUniform(), μ, x)
+    return _transport_from_mvstd(ν, StdUniform(), z)
+end
+
 @generated function _transport_between_origins(
     ν,
     ::StaticInteger{n_ν},
diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl
index a1d8865a..1503f6e2 100644
--- a/test/combinators/bind.jl
+++ b/test/combinators/bind.jl
@@ -64,6 +64,23 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca
         @test all(u -> 0 <= u <= 1, y)
         xy_reco = transport_to(μ, StdUniform()^3)(y)
         @test xy_reco ≈ xy
+
+        # Transport between two measures of unknown DOF (mvstd pivot):
+        μ2 = mbind(f_βv, StdUniform()^1, vcat)
+        xy2 = transport_to(μ2, μ)(xy)
+        @test xy2 isa AbstractVector{<:Real} && length(xy2) == 3
+        @test transport_to(μ, μ2)(xy2) ≈ xy
+        @test logdensityof(μ2, xy2) isa Real
+
+        # Transport between known-DOF and unknown-DOF measures (mvstd pivot):
+        ν_known = productmeasure((StdNormal(), StdNormal(), StdNormal()))
+        z = transport_to(ν_known, μ)(xy)
+        @test z isa Tuple{Vararg{Real,3}}
+        @test collect(transport_to(μ, ν_known)(z)) ≈ xy
+
+        # DOF mismatches must not go unnoticed:
+        @test_throws ArgumentError transport_to(StdUniform()^5, μ)(xy)
+        @test_throws ArgumentError transport_to(StdUniform()^2, μ)(xy)
     end
 
     @testset "mbind with merge" begin

From a13275b8b9e338445a61c316b8cf17f3b23e74af Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 21:58:48 +0200
Subject: [PATCH 055/122] Make superpose simplifications type stable

Superposition simplifications now only happen when measure equality is
decidable from the measure types alone (via the new _static_isequal),
instead of branching on runtime measure equality. As a consequence,
value-equal measures of non-singleton type (e.g. equal Diracs) now stay
superpositions instead of collapsing to weighted measures, and
basemeasure of superpositions with equal-typed components is type stable.

Created by generative AI.
---
 src/combinators/smart-constructors.jl  | 30 ++++++++++++++++++++------
 test/combinators/smart_constructors.jl | 12 +++++++++++
 test/combinators/superpose.jl          | 13 ++++++++---
 3 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index db355184..e2c4cfe3 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -184,6 +184,8 @@ as a collection (array, tuple or named tuple) of measures.
 The vararg form simplifies algebraically: equal measures combine into weighted
 measures (`superpose(μ, μ) == weightedmeasure(log(2), μ)`), weighted measures
 with equal bases add their weights, and superpositions merge their components.
+To keep `superpose` type stable, such simplifications only happen when
+equality of the measures involved can be decided from their types alone.
 Collections are wrapped as-is, apart from cost-free structural simplifications.
 """
 function superpose end
@@ -208,12 +210,22 @@ superpose(a::FillArrays.Fill) = weightedmeasure(log(length(a)), asmeasure(_fill_
 superpose(t::Tuple) = SuperpositionMeasure(t)
 superpose(nt::NamedTuple) = SuperpositionMeasure(nt)
 
+# Measure equality can typically only be established at runtime, but measure
+# construction must be type stable, so simplifications may only depend on
+# measure equality that is decidable from the measure types alone:
+@inline _static_isequal(::T, ::T) where {T} = static(Base.issingletontype(T))
+@inline _static_isequal(::Any, ::Any) = static(false)
+
 function _superpose_two(μ::AbstractMeasure, ν::AbstractMeasure)
-    μ == ν ? weightedmeasure(static(float(logtwo)), μ) : SuperpositionMeasure((μ, ν))
+    if _static_isequal(μ, ν) isa True
+        weightedmeasure(static(float(logtwo)), μ)
+    else
+        SuperpositionMeasure((μ, ν))
+    end
 end
 
 function _superpose_two(μ::WeightedMeasure, ν::WeightedMeasure)
-    if μ.base == ν.base
+    if _static_isequal(μ.base, ν.base) isa True
         weightedmeasure(logaddexp(asnonstatic(μ.logweight), asnonstatic(ν.logweight)), μ.base)
     else
         SuperpositionMeasure((μ, ν))
@@ -221,13 +233,19 @@ function _superpose_two(μ::WeightedMeasure, ν::WeightedMeasure)
 end
 
 function _superpose_two(μ::WeightedMeasure, ν::AbstractMeasure)
-    μ.base == ν ? weightedmeasure(log1pexp(asnonstatic(μ.logweight)), μ.base) :
-    SuperpositionMeasure((μ, ν))
+    if _static_isequal(μ.base, ν) isa True
+        weightedmeasure(log1pexp(asnonstatic(μ.logweight)), μ.base)
+    else
+        SuperpositionMeasure((μ, ν))
+    end
 end
 
 function _superpose_two(μ::AbstractMeasure, ν::WeightedMeasure)
-    μ == ν.base ? weightedmeasure(log1pexp(asnonstatic(ν.logweight)), ν.base) :
-    SuperpositionMeasure((μ, ν))
+    if _static_isequal(μ, ν.base) isa True
+        weightedmeasure(log1pexp(asnonstatic(ν.logweight)), ν.base)
+    else
+        SuperpositionMeasure((μ, ν))
+    end
 end
 
 _superpose_two(μ::SuperpositionMeasure, ν::SuperpositionMeasure) =
diff --git a/test/combinators/smart_constructors.jl b/test/combinators/smart_constructors.jl
index ac6da161..06d7906b 100644
--- a/test/combinators/smart_constructors.jl
+++ b/test/combinators/smart_constructors.jl
@@ -103,6 +103,18 @@ using Static: static
         sv2 = superpose(sv, StdExponential())
         @test length(sv.components) == 2 && length(sv2.components) == 3
 
+        # Simplifications must be type stable, so they only happen when
+        # measure equality is decidable from the measure types:
+        @inferred superpose(μ, ν)
+        @inferred superpose(μ, μ)
+        @inferred superpose(2.0 * μ, 3.0 * μ)
+        @inferred superpose(2.0 * μ, μ)
+        @inferred superpose(μ, μ, μ, μ)
+        @inferred superpose(Dirac(1), Dirac(1))
+        @test superpose(Dirac(1), Dirac(1)) isa SuperpositionMeasure
+        @inferred superpose(2.0 * Dirac(1), 3.0 * Dirac(1))
+        @test superpose(2.0 * Dirac(1), 3.0 * Dirac(1)) isa SuperpositionMeasure
+
         @test superpose(Fill(μ, 4)) == weightedmeasure(log(4), μ)
         @test superpose([μ, μ, μ]) == weightedmeasure(log(3), μ)
 
diff --git a/test/combinators/superpose.jl b/test/combinators/superpose.jl
index 17c677df..753d5b8c 100644
--- a/test/combinators/superpose.jl
+++ b/test/combinators/superpose.jl
@@ -1,7 +1,7 @@
 using Test
 
 using MeasureBase
-using MeasureBase: superpose
+using MeasureBase: superpose, weightedmeasure, StdNormal
 
 @testset "superpose.jl" begin
     μ = Dirac(0)
@@ -17,8 +17,15 @@ using MeasureBase: superpose
     @test density_def(μs, 0) == 1.0
     @test basemeasure(μs) == weightedmeasure(log(2), CountingBase())
 
+    # Dirac equality is not decidable from types, so no weighted collapse:
     μ2 = μ + μ
-    @test μ2 isa WeightedMeasure
+    @test μ2 isa SuperpositionMeasure
     @test μ2 == superpose(μ, μ)
-    @test basemeasure(μ2) == μ
+    @test density_def(μ2, 0) == 1.0
+
+    # For singleton measure types equal measures combine into weighted measures:
+    s2 = StdNormal() + StdNormal()
+    @test s2 isa WeightedMeasure
+    @test exp(s2.logweight) ≈ 2
+    @test basemeasure(s2) == StdNormal()
 end

From 2068175a50a51d94bbe262323c73c21f97c341b7 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 22:06:55 +0200
Subject: [PATCH 056/122] Rework non-relative density evaluation around
 logdensityof_impl and with-rest

logdensityof is now a single generic entry function; measure types
specialize the new MeasureBase.logdensityof_impl instead of logdensityof
itself. The new logdensityof_with_rest protocol evaluates densities of
measures that live at the beginning of a flat variate stream (a vector
for vcat-combined and a NamedTuple for merge-combined measures),
returning the log-density, the consumed variate and the unconsumed rest
of the stream. Its default implementation determines the variate size
via mspace_elsize (falling back to testvalue).

Bind and CombinedMeasure evaluate densities of vcat- and merge-combined
variates through logdensityof_with_rest in a single pass now, without
materializing intermediate transport measures and without splitting
sub-variates twice. Variates that are too long now result in an
informative exception, and so do binds whose value combination function
does not support variate splitting.

Created by generative AI.
---
 src/collection_utils.jl               |  9 ++++
 src/combinators/bind.jl               | 53 ++++++++++++++++----
 src/combinators/combined.jl           | 47 +++++++++++++++---
 src/combinators/half.jl               |  2 +-
 src/combinators/power.jl              |  8 ++--
 src/combinators/product.jl            | 10 ++--
 src/combinators/transformedmeasure.jl | 10 ++--
 src/density-core.jl                   | 69 +++++++++++++++++++++++----
 src/density.jl                        |  2 +-
 src/primitive.jl                      |  4 +-
 src/primitives/counting.jl            |  4 +-
 src/primitives/dirac.jl               |  4 +-
 src/primitives/lebesgue.jl            |  4 +-
 src/standard/stdexponential.jl        |  2 +-
 src/standard/stdlogistic.jl           |  2 +-
 src/standard/stdnormal.jl             |  2 +-
 src/standard/stduniform.jl            |  2 +-
 test/combinators/bind.jl              | 10 ++++
 18 files changed, 192 insertions(+), 52 deletions(-)

diff --git a/src/collection_utils.jl b/src/collection_utils.jl
index be5cc934..6b9de0d5 100644
--- a/src/collection_utils.jl
+++ b/src/collection_utils.jl
@@ -101,3 +101,12 @@ _cat_measures(a::Tuple, b::Tuple) = (a..., b...)
 _cat_measures(a::AbstractVector, b::Tuple) = vcat(a, [b...])
 _cat_measures(a::Tuple, b::AbstractVector) = vcat([a...], b)
 _cat_measures(a::AbstractVector, b::AbstractVector) = vcat(a, b)
+
+
+# Take the beginning of a flat vector stream as a variate of size `sz`:
+Base.@propagate_inbounds _consume_from_stream(x::AbstractVector, sz::Tuple{IntegerLike}) =
+    _split_after(x, sz[1])
+
+function _consume_from_stream(x::AbstractVector, @nospecialize(sz::Tuple))
+    throw(ArgumentError("Can't consume a variate of size $sz from a flat vector stream"))
+end
diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl
index 27cb02a4..ab7ab708 100644
--- a/src/combinators/bind.jl
+++ b/src/combinators/bind.jl
@@ -284,18 +284,51 @@ testvalue(::Bind) = throw(ArgumentError("testvalue is not available for Bind"))
 logdensity_def(::Bind, x) =
     throw(ArgumentError("logdensity_def is not available for Bind"))
 
-# Specialize logdensityof to avoid duplicate calculations:
-function logdensityof(μ::Bind, x)
-    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, x)
-    β_a = _get_β_a(μ, a)
-    logdensityof(tpm_α, a) + logdensityof(β_a, b)
+# Density evaluation consumes the variate parts of the primary and secondary
+# measure in a single pass, using the with-rest protocol for value-dependent
+# variate sizes:
+
+logdensityof_impl(μ::Bind, x) = _bind_ld_impl(μ.f_c, μ, x)
+
+unsafe_logdensityof(μ::Bind, x) = logdensityof_impl(μ, x)
+
+function _bind_ld_impl(::typeof(tuple), μ::Bind, xy::Tuple{Vararg{Any,2}})
+    a, b = xy[1], xy[2]
+    logdensityof(μ.α, a) + logdensityof(_get_β_a(μ, a), b)
 end
 
-# Specialize unsafe_logdensityof to avoid duplicate calculations:
-function unsafe_logdensityof(μ::Bind, x)
-    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, x)
-    β_a = _get_β_a(μ, a)
-    unsafe_logdensityof(tpm_α, a) + unsafe_logdensityof(β_a, b)
+function _bind_ld_impl(::Type{Pair}, μ::Bind, xy::Pair)
+    a, b = xy.first, xy.second
+    logdensityof(μ.α, a) + logdensityof(_get_β_a(μ, a), b)
+end
+
+function _bind_ld_impl(::Union{typeof(vcat),typeof(merge)}, μ::Bind, xy)
+    ℓ, x_μ, x_rest = logdensityof_with_rest(μ, xy)
+    if !isempty(x_rest)
+        throw(ArgumentError("Variate too long during density evaluation of a bind"))
+    end
+    return ℓ
+end
+
+function _bind_ld_impl(@nospecialize(f_c), @nospecialize(μ::Bind), @nospecialize(xy))
+    throw(
+        ArgumentError(
+            "Can't compute density of a bind with value combination function of type $(nameof(typeof(f_c)))",
+        ),
+    )
+end
+
+function logdensityof_with_rest(μ::_BindBy{typeof(vcat)}, x::AbstractVector)
+    ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x)
+    ℓ_b, b, x_rest = logdensityof_with_rest(_get_β_a(μ, a), x2)
+    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    return ℓ_a + ℓ_b, x_μ, x_rest
+end
+
+function logdensityof_with_rest(μ::_BindBy{typeof(merge)}, x::NamedTuple)
+    ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x)
+    ℓ_b, b, x_rest = logdensityof_with_rest(_get_β_a(μ, a), x2)
+    return ℓ_a + ℓ_b, merge(a, b), x_rest
 end
 
 
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index c3ea079c..4325ced8 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -134,15 +134,50 @@ function logdensity_def(μ::CombinedMeasure, ab)
     return logdensity_def(tpm_α, a) + logdensity_def(μ.β, b)
 end
 
-# Specialize logdensityof directly to avoid creating temporary combined base measures:
-function logdensityof(μ::CombinedMeasure, ab)
-    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab)
+# Density evaluation consumes the variate parts of both component measures
+# in a single pass, using the with-rest protocol for value-dependent
+# variate sizes:
+
+logdensityof_impl(μ::CombinedMeasure, ab) = _combined_ld_impl(μ.f_c, μ, ab)
+
+unsafe_logdensityof(μ::CombinedMeasure, ab) = logdensityof_impl(μ, ab)
+
+function _combined_ld_impl(::typeof(tuple), μ::CombinedMeasure, ab::Tuple{Vararg{Any,2}})
+    logdensityof(μ.α, ab[1]) + logdensityof(μ.β, ab[2])
+end
+
+function _combined_ld_impl(::Type{Pair}, μ::CombinedMeasure, ab::Pair)
+    logdensityof(μ.α, ab.first) + logdensityof(μ.β, ab.second)
+end
+
+function _combined_ld_impl(::Union{typeof(vcat),typeof(merge)}, μ::CombinedMeasure, ab)
+    ℓ, x_μ, x_rest = logdensityof_with_rest(μ, ab)
+    if !isempty(x_rest)
+        throw(
+            ArgumentError(
+                "Variate too long during density evaluation of a combined measure",
+            ),
+        )
+    end
+    return ℓ
+end
+
+function _combined_ld_impl(f_c, μ::CombinedMeasure, ab)
+    tpm_α, a, b = tpmeasure_split_combined(f_c, μ.α, ab)
     return logdensityof(tpm_α, a) + logdensityof(μ.β, b)
 end
 
-function unsafe_logdensityof(μ::CombinedMeasure, ab)
-    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab)
-    return unsafe_logdensityof(tpm_α, a) + unsafe_logdensityof(μ.β, b)
+function logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector)
+    ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x)
+    ℓ_b, b, x_rest = logdensityof_with_rest(μ.β, x2)
+    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    return ℓ_a + ℓ_b, x_μ, x_rest
+end
+
+function logdensityof_with_rest(μ::CombinedMeasure{typeof(merge)}, x::NamedTuple)
+    ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x)
+    ℓ_b, b, x_rest = logdensityof_with_rest(μ.β, x2)
+    return ℓ_a + ℓ_b, merge(a, b), x_rest
 end
 
 
diff --git a/src/combinators/half.jl b/src/combinators/half.jl
index 24063b76..c93327b0 100644
--- a/src/combinators/half.jl
+++ b/src/combinators/half.jl
@@ -19,7 +19,7 @@ function Base.rand(rng::AbstractRNG, ::Type{T}, μ::Half) where {T}
     return abs(rand(rng, T, unhalf(μ)))
 end
 
-function logdensityof(μ::Half, x)
+function logdensityof_impl(μ::Half, x)
     ld = logdensityof(unhalf(μ), x) - loghalf
     return x ≥ 0 ? ld : oftype(ld, -Inf)
 end
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 1ce9ca98..5145a80c 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -97,8 +97,8 @@ params(d::PowerMeasure) = params(first(marginals(d)))
     basemeasure(d.parent)^d.axes
 end
 
-for func in [:logdensityof, :logdensity_def]
-    @eval @inline function $func(d::PowerMeasure{M}, x) where {M}
+for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)]
+    @eval @inline function $head(d::PowerMeasure{M}, x) where {M}
         parent_m = d.parent
         sz_parent = axes2size(d.axes)
         sz_x = maybestatic_size(x)
@@ -114,7 +114,7 @@ for func in [:logdensityof, :logdensity_def]
         end
     end
 
-    @eval @inline function $func(
+    @eval @inline function $head(
         d::PowerMeasure{<:Any,Tuple{<:StaticOneToLike{N}}},
         x,
     ) where {N}
@@ -124,7 +124,7 @@ for func in [:logdensityof, :logdensity_def]
         end
     end
 
-    @eval @inline function $func(
+    @eval @inline function $head(
         ::PowerMeasure{<:Any,<:Tuple{Vararg{StaticOneToLike{0}}}},
         x,
     )
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 656ded0d..9899e0a9 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -72,8 +72,8 @@ function _rand_product(
     end |> collect
 end
 
-for func in [:logdensityof, :logdensity_def]
-    @eval @inline function $func(d::AbstractProductMeasure, x)
+for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)]
+    @eval @inline function $head(d::AbstractProductMeasure, x)
         mapreduce($func, +, marginals(d), x)
     end
 end
@@ -112,14 +112,14 @@ end
     return q
 end
 
-for func in [:logdensityof, :logdensity_def]
+for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)]
     # For tuples, `mapreduce` has trouble with type inference
-    @eval @inline function $func(d::ProductMeasure{T}, x) where {T<:Tuple}
+    @eval @inline function $head(d::ProductMeasure{T}, x) where {T<:Tuple}
         ℓs = map($func, marginals(d), x)
         sum(ℓs)
     end
 
-    @eval function $func(d::ProductMeasure{NamedTuple{N,T}}, x) where {N,T}
+    @eval function $head(d::ProductMeasure{NamedTuple{N,T}}, x) where {N,T}
         _product_gen_impl(Val($func), d, x)
     end
 end
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 0017c86f..ee960f9a 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -135,7 +135,7 @@ function _combine_logd_with_ladj(logd_orig::Real, ladj::Real)
     end
 end
 
-function logdensityof(
+function logdensityof_impl(
     @nospecialize(μ::_NonBijectivePusfwdMeasure{M,<:PushfwdRootMeasure}),
     @nospecialize(v::Any)
 ) where {M}
@@ -146,7 +146,7 @@ function logdensityof(
     )
 end
 
-function logdensityof(
+function logdensityof_impl(
     @nospecialize(μ::_NonBijectivePusfwdMeasure{M,<:AdaptRootMeasure}),
     @nospecialize(v::Any)
 ) where {M}
@@ -157,15 +157,15 @@ function logdensityof(
     )
 end
 
-for func in [:logdensityof, :logdensity_def]
-    @eval function $func(ν::PushforwardMeasure{F,I,M,<:AdaptRootMeasure}, y) where {F,I,M}
+for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)]
+    @eval function $head(ν::PushforwardMeasure{F,I,M,<:AdaptRootMeasure}, y) where {F,I,M}
         f_inv = unwrap(ν.finv)
         x, inv_ladj = with_logabsdet_jacobian(f_inv, y)
         logd_orig = $func(ν.origin, x)
         return _combine_logd_with_ladj(logd_orig, inv_ladj)
     end
 
-    @eval function $func(ν::PushforwardMeasure{F,I,M,<:PushfwdRootMeasure}, y) where {F,I,M}
+    @eval function $head(ν::PushforwardMeasure{F,I,M,<:PushfwdRootMeasure}, y) where {F,I,M}
         f_inv = unwrap(ν.finv)
         x = f_inv(y)
         logd_orig = $func(ν.origin, x)
diff --git a/src/density-core.jl b/src/density-core.jl
index 33a3c3cb..806a33a6 100644
--- a/src/density-core.jl
+++ b/src/density-core.jl
@@ -10,25 +10,41 @@ export density_rel
 export density_def
 
 """
-    logdensityof(m::AbstractMeasure, x) 
+    logdensityof(m::AbstractMeasure, x)
 
 Compute the log-density of the measure `m` at `x`. Density is always relative,
 but `DensityInterface.jl` does not account for this. For compatibility with
 this, `logdensityof` for a measure is always implicitly relative to
-[`rootmeasure(x)`](@ref rootmeasure). 
+[`rootmeasure(x)`](@ref rootmeasure).
 
-`logdensityof` works by first computing `insupport(m, x)`. If this is true, then
-`unsafe_logdensityof` is called. If `insupport(m, x)` is known to be `true`, it
-can be a little faster to directly call `unsafe_logdensityof(m, x)`. 
+`logdensityof(m, x)` is implemented via
+[`MeasureBase.logdensityof_impl`](@ref), measure types should specialize
+`logdensityof_impl` instead of `logdensityof` itself.
 
 To compute log-density relative to `basemeasure(m)` or *define* a log-density
 (relative to `basemeasure(m)` or another measure given explicitly), see
-`logdensity_def`. 
+`logdensity_def`.
 
 To compute a log-density relative to a specific base-measure, see
-`logdensity_rel`. 
+`logdensity_rel`.
 """
-@inline function logdensityof(μ::AbstractMeasure, x)
+@inline logdensityof(μ::AbstractMeasure, x) = logdensityof_impl(μ, x)
+
+"""
+    MeasureBase.logdensityof_impl(μ::AbstractMeasure, x)
+
+Implements [`logdensityof(μ, x)`](@ref logdensityof).
+
+Measure types should specialize `logdensityof_impl` instead of
+`logdensityof` itself. Implementations must return the log-density of `μ`
+at `x` relative to [`rootmeasure(μ)`](@ref) and must handle `x` outside of
+the support of `μ` (the result must be `-Inf` then).
+
+The default implementation checks `insupport(μ, x)` (unless the result is
+a [`MeasureBase.NoFastInsupport`](@ref)) and computes the density via
+[`unsafe_logdensityof`](@ref).
+"""
+@inline function logdensityof_impl(μ::AbstractMeasure, x)
     result = dynamic(unsafe_logdensityof(μ, x))
     _checksupport(insupport(μ, x), result)
 end
@@ -40,6 +56,43 @@ end
 _checksupport(cond, result) = ifelse(cond == true, result, oftype(result, -Inf))
 @inline _checksupport(::NoFastInsupport, result) = result
 
+"""
+    MeasureBase.logdensityof_with_rest(μ::AbstractMeasure, x)
+
+Compute the log-density of `μ` at the beginning of `x`, a flat stream of
+variate content that may extend beyond the variate of `μ`.
+
+`x` must either be a vector that starts with the (flattened) variate of `μ`
+(for measures combined via `vcat`) or a `NamedTuple` whose first properties
+constitute the variate of `μ` (for measures combined via `merge`).
+
+Returns a tuple `(ℓ, x_μ, x_rest)` of the log-density `ℓ`, the variate
+`x_μ` of `μ` consumed from the stream, and the unconsumed rest of the
+stream.
+
+Measure types whose variate size depends on measure values, like
+[`mbind`](@ref) results, implement density calculation via
+`logdensityof_with_rest` instead of
+[`logdensityof_impl`](@ref MeasureBase.logdensityof_impl).
+
+The default implementation determines the size resp. the property names of
+the variate via [`some_mspace_elsize`](@ref MeasureBase.some_mspace_elsize)
+resp. `testvalue` and delegates to `logdensityof_impl`.
+"""
+function logdensityof_with_rest end
+
+function logdensityof_with_rest(μ::AbstractMeasure, x::AbstractVector)
+    a, x_rest = _consume_from_stream(x, some_mspace_elsize(μ))
+    return logdensityof_impl(μ, a), a, x_rest
+end
+
+function logdensityof_with_rest(μ::AbstractMeasure, x::NamedTuple)
+    a, x_rest = _split_after(x, Val(_mspace_names(μ)))
+    return logdensityof_impl(μ, a), a, x_rest
+end
+
+_mspace_names(μ::AbstractMeasure) = keys(testvalue(μ))
+
 
 """
     localmeasure(m::AbstractMeasure, x)::AbstractMeasure
diff --git a/src/density.jl b/src/density.jl
index 06dc98e1..67d34464 100644
--- a/src/density.jl
+++ b/src/density.jl
@@ -224,7 +224,7 @@ logdensity_def(μ::DensityMeasure, x) = logdensityof(μ.f, x)
 
 density_def(μ::DensityMeasure, x) = densityof(μ.f, x)
 
-function logdensityof(μ::DensityMeasure, x::Any)
+function logdensityof_impl(μ::DensityMeasure, x::Any)
     integrand, μ_base = μ.f, μ.base
 
     base_logval = logdensityof(μ_base, x)
diff --git a/src/primitive.jl b/src/primitive.jl
index 85cf2beb..f43485c2 100644
--- a/src/primitive.jl
+++ b/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 logdensityof_impl(::PrimitiveMeasure, x::Real) = zero(float(typeof(x)))
+@inline logdensityof_impl(::PrimitiveMeasure, x) = false
 
 logdensity_def(::PrimitiveMeasure, x) = static(0.0)
 
diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl
index 31398f32..5a7a998e 100644
--- a/src/primitives/counting.jl
+++ b/src/primitives/counting.jl
@@ -12,12 +12,12 @@ struct Counting{T} <: AbstractMeasure
     Counting(supp) = new{Core.Typeof(supp)}(supp)
 end
 
-@inline function logdensityof(μ::Counting, x::Real)
+@inline function logdensityof_impl(μ::Counting, x::Real)
     R = float(typeof(x))
     insupport(μ, x) ? zero(R) : R(-Inf)
 end
 
-@inline logdensityof(μ::Counting, x) = insupport(μ, x) ? 0.0 : -Inf
+@inline logdensityof_impl(μ::Counting, x) = insupport(μ, x) ? 0.0 : -Inf
 
 @inline logdensity_def(μ::Counting, x) = logdensityof(μ, x)
 
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index f5e5931a..ba044c1d 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -23,12 +23,12 @@ basemeasure(d::Dirac) = CountingBase()
 
 massof(::Dirac) = static(1.0)
 
-function logdensityof(μ::Dirac, x::Real)
+function logdensityof_impl(μ::Dirac, x::Real)
     R = float(typeof(x))
     insupport(μ, x) ? zero(R) : R(-Inf)
 end
 
-logdensityof(μ::Dirac, x) = insupport(μ, x) ? 0.0 : -Inf
+logdensityof_impl(μ::Dirac, x) = insupport(μ, x) ? 0.0 : -Inf
 
 logdensity_def(::Dirac, x::Real) = zero(float(typeof(x)))
 logdensity_def(::Dirac, x) = 0.0
diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl
index 4b8bf7ab..0e6b171a 100644
--- a/src/primitives/lebesgue.jl
+++ b/src/primitives/lebesgue.jl
@@ -63,12 +63,12 @@ insupport(μ::Lebesgue, x) = x ∈ μ.support
 
 insupport(::Lebesgue{RealValues}, ::Real) = true
 
-@inline function logdensityof(μ::Lebesgue, x::Real)
+@inline function logdensityof_impl(μ::Lebesgue, x::Real)
     R = float(typeof(x))
     insupport(μ, x) ? zero(R) : R(-Inf)
 end
 
-@inline logdensityof(μ::Lebesgue, x) = insupport(μ, x) ? 0.0 : -Inf
+@inline logdensityof_impl(μ::Lebesgue, x) = insupport(μ, x) ? 0.0 : -Inf
 
 massof(::Lebesgue{RealValues}, s::Interval) = width(s)
 
diff --git a/src/standard/stdexponential.jl b/src/standard/stdexponential.jl
index c985c224..fda5b53a 100644
--- a/src/standard/stdexponential.jl
+++ b/src/standard/stdexponential.jl
@@ -4,7 +4,7 @@ export StdExponential
 
 insupport(::StdExponential, x) = x ≥ zero(x)
 
-@inline function logdensityof(::StdExponential, x)
+@inline function logdensityof_impl(::StdExponential, x)
     R = float(typeof(x))
     x ≥ zero(R) ? convert(R, -x) : R(-Inf)
 end
diff --git a/src/standard/stdlogistic.jl b/src/standard/stdlogistic.jl
index 58a1ba67..b28dd618 100644
--- a/src/standard/stdlogistic.jl
+++ b/src/standard/stdlogistic.jl
@@ -4,7 +4,7 @@ export StdLogistic
 
 @inline insupport(d::StdLogistic, x) = true
 
-@inline logdensityof(::StdLogistic, x) = (u = -abs(x); u - 2 * log1pexp(u))
+@inline logdensityof_impl(::StdLogistic, x) = (u = -abs(x); u - 2 * log1pexp(u))
 
 @inline logdensity_def(::StdLogistic, x) = logdensityof(StdLogistic(), x)
 @inline basemeasure(::StdLogistic) = LebesgueBase()
diff --git a/src/standard/stdnormal.jl b/src/standard/stdnormal.jl
index 057a8629..f636d311 100644
--- a/src/standard/stdnormal.jl
+++ b/src/standard/stdnormal.jl
@@ -7,7 +7,7 @@ export StdNormal
 
 @inline insupport(::StdNormal, x) = true
 
-@inline logdensityof(::StdNormal, x) = (-x^2 - log2π) / 2
+@inline logdensityof_impl(::StdNormal, x) = (-x^2 - log2π) / 2
 
 @inline logdensity_def(::StdNormal, x) = -x^2 / 2
 @inline basemeasure(::StdNormal) = WeightedMeasure(static(-0.5 * log2π), LebesgueBase())
diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl
index e3702656..d0fc236d 100644
--- a/src/standard/stduniform.jl
+++ b/src/standard/stduniform.jl
@@ -4,7 +4,7 @@ export StdUniform
 
 insupport(::StdUniform, x) = zero(x) ≤ x ≤ one(x)
 
-@inline function logdensityof(::StdUniform, x)
+@inline function logdensityof_impl(::StdUniform, x)
     R = float(typeof(x))
     zero(x) ≤ x ≤ one(x) ? zero(R) : R(-Inf)
 end
diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl
index 1503f6e2..f7621037 100644
--- a/test/combinators/bind.jl
+++ b/test/combinators/bind.jl
@@ -81,6 +81,16 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca
         # DOF mismatches must not go unnoticed:
         @test_throws ArgumentError transport_to(StdUniform()^5, μ)(xy)
         @test_throws ArgumentError transport_to(StdUniform()^2, μ)(xy)
+
+        # Nested binds evaluate in a single with-rest pass:
+        μnest = mbind(f_βv, μ, vcat)
+        xyz = rand(stblrng(), Float64, μnest)
+        @test length(xyz) == 5
+        an, bn = xyz[1:3], xyz[4:5]
+        @test logdensityof(μnest, xyz) ≈ logdensityof(μ, an) + logdensityof(f_βv(an), bn)
+
+        # Variates that are too long must not go unnoticed:
+        @test_throws ArgumentError logdensityof(μ, vcat(xy, [0.5]))
     end
 
     @testset "mbind with merge" begin

From e4bdb7b1fe47fcbb976d3470e46b60366119e7a5 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 22:23:58 +0200
Subject: [PATCH 057/122] Rework relative density evaluation as type-stable
 lockstep chain descent

The generic three-argument logdensity_def now descends the base measure
chains of both measures in lockstep after equalizing their static depths.
Since the members of a shared chain suffix have the same depth-from-root
on both sides, shared suffixes cancel symbolically. The descent is fully
unrolled at compile time and constructs only the base measures it
actually visits.

Specialized relative densities move from three-argument logdensity_def
methods to the new extension point MeasureBase.logdensity_rel_def. The
descent checks for an applicable specialization at each visited measure
pair; availability is decided purely by dispatch via a sentinel return
type instead of method-table introspection, so the checks are free at
run time and defining new specializations behaves like any ordinary
method definition.

This removes basemeasure_sequence-based chain materialization, the
commonbase search, schema and the static_hasmethod gate from the
relative density code path. Root measure pairs without a specialization
now throw an informative exception instead of warning and returning NaN.

Created by generative AI.
---
 src/MeasureBase.jl           |   8 +-
 src/combinators/superpose.jl |   8 +-
 src/density-core.jl          | 163 ++++++++++++++++++++++-------------
 src/interface.jl             |   3 +-
 src/primitive.jl             |   4 +-
 src/primitives/lebesgue.jl   |   4 +-
 src/schema.jl                |  34 --------
 src/utils.jl                 |  30 -------
 test/test_basics.jl          |  23 +++++
 9 files changed, 142 insertions(+), 135 deletions(-)
 delete mode 100644 src/schema.jl

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 9f8037eb..e4a3ce42 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -162,6 +162,13 @@ Compute the log-density of the measure m at the point `x`, relative to
 
 Compute the log-density of `m1` relative to `m2` at the point `x`, assuming
 `insupport(m1, x)` and `insupport(m2, x)`.
+
+The generic implementation descends the base measure chains of both
+measures in lockstep, so it terminates at the first specialized
+`logdensity_def` method for a pair of base measures (in particular at pairs
+of identical primitive measures) and any shared chain suffix cancels
+symbolically. Measure types may add specialized three-argument methods for
+measure pairs whose relative density can be computed more directly.
 """
 function logdensity_def end
 
@@ -175,7 +182,6 @@ include("smf.jl")
 include("mspace.jl")
 include("getdof.jl")
 include("transport.jl")
-include("schema.jl")
 include("proxies.jl")
 include("kernel.jl")
 include("parameterized.jl")
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index aa7b6e20..d9b53bd5 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -82,7 +82,7 @@ function density_def(s::SuperpositionMeasure, x)
     end
 end
 
-@inline function logdensity_def(μ::T, ν::T, x) where {T<:SuperpositionMeasure}
+@inline function logdensity_rel_def(μ::T, ν::T, x) where {T<:SuperpositionMeasure}
     if μ === ν
         return zero(return_type(logdensity_def, (μ, x)))
     else
@@ -98,12 +98,12 @@ function _superpos_logdensity_rel(s::SuperpositionMeasure, β, x)
     logsumexp(ds)
 end
 
-@inline logdensity_def(s::SuperpositionMeasure, β, x) = _superpos_logdensity_rel(s, β, x)
+@inline logdensity_rel_def(s::SuperpositionMeasure, β, x) = _superpos_logdensity_rel(s, β, x)
 
-@inline logdensity_def(s::SuperpositionMeasure, β::SuperpositionMeasure, x) =
+@inline logdensity_rel_def(s::SuperpositionMeasure, β::SuperpositionMeasure, x) =
     _superpos_logdensity_rel(s, β, x)
 
-@inline logdensity_def(s, β::SuperpositionMeasure, x) = -_superpos_logdensity_rel(β, s, x)
+@inline logdensity_rel_def(s, β::SuperpositionMeasure, x) = -_superpos_logdensity_rel(β, s, x)
 
 @inline logdensity_def(s::SuperpositionMeasure, x) = log(density_def(s, x))
 
diff --git a/src/density-core.jl b/src/density-core.jl
index 806a33a6..9a6b0649 100644
--- a/src/density-core.jl
+++ b/src/density-core.jl
@@ -230,75 +230,118 @@ See also `logdensity_rel`.
 @inline function unsafe_logdensity_rel(μ::AbstractMeasure, ν::AbstractMeasure, x)
     μ_local = localmeasure(μ, x)
     ν_local = localmeasure(ν, x)
-    # Extra dispatch boundary to reduce number of required specializations of implementation:
-    return _unsafe_logdensity_rel_local(μ_local, ν_local, x)
+    return logdensity_def(μ_local, ν_local, x)
 end
 
-@inline function _unsafe_logdensity_rel_local(μ::M, ν::N, x::X) where {M,N,X}
-    if static_hasmethod(logdensity_def, Tuple{M,N,X})
-        return logdensity_def(μ, ν, x)
-    end
-    μs = basemeasure_sequence(μ)
-    νs = basemeasure_sequence(ν)
-    cb = commonbase(μs, νs, X)
-    # _logdensity_rel(μ, ν)
-    isnothing(cb) && begin
-        μ = μs[end]
-        ν = νs[end]
-        @warn """
-        No common base measure for
-            $μ
-        and
-            $ν
-
-        Returning a relative log-density of NaN. If this is incorrect, add a
-        three-argument method
-            logdensity_def($μ, $ν, x)
-        """
-        return NaN
-    end
-    return _logdensity_rel(μs, νs, cb, x)
-end
+# Indicates that no specialized method is available to compute the
+# log-density between a given pair of measures:
+struct _NoLogdensityRel end
 
-# Note that this method assumes `μ` and `ν` to have the same type
-function logdensity_def(μ::T, ν::T, x) where {T}
-    if μ === ν
-        return zero(logdensity_def(μ, x))
-    else
-        α = basemeasure(μ)
-        β = basemeasure(ν)
-        return logdensity_def(μ, x) - logdensity_def(ν, x) + logdensity_rel(α, β, x)
-    end
-end
+"""
+    MeasureBase.logdensity_rel_def(μ, ν, x)
 
-@generated function _logdensity_rel(
-    μs::Tμ,
-    νs::Tν,
-    ::Tuple{<:StaticInteger{M},<:StaticInteger{N}},
-    x::X,
-) where {Tμ,Tν,M,N,X}
-    sμ = schema(Tμ)
-    sν = schema(Tν)
+Specialization point for the log-density of `μ` relative to `ν` at `x`.
 
-    q = quote
-        $(Expr(:meta, :inline))
-        ℓ = logdensity_def(μs[$M], νs[$N], x)
-    end
+Measure types may add methods for pairs of measure types whose relative
+density can be computed directly. The generic implementation of
+[`logdensity_def(μ, ν, x)`](@ref logdensity_def) descends the base measure
+chains of both measures in lockstep and uses the first specialized
+`logdensity_rel_def` method it encounters along the way.
 
-    for i in 1:(M-1)
-        push!(q.args, :(Δℓ = logdensity_def(μs[$i], x)))
-        # push!(q.args, :(println("Adding", Δℓ)))
-        push!(q.args, :(ℓ += Δℓ))
-    end
+Do not call `logdensity_rel_def` directly, call
+[`logdensity_rel`](@ref) (or `logdensity_def`) instead.
+"""
+@inline logdensity_rel_def(μ, ν, x) = _NoLogdensityRel()
+
+# Generic relative density: descend the base measure chains of both measures
+# in lockstep, after equalizing their depths. Since the members of a shared
+# chain suffix have the same depth-from-root on both sides, the descent
+# terminates at a specialized `logdensity_rel_def` method as soon as one
+# becomes applicable (in particular for pairs of identical primitive
+# measures), so any shared chain suffix cancels symbolically instead of
+# numerically. The descent is fully unrolled at compile time based on the
+# static base measure depths, only the base measures actually visited are
+# constructed, and whether a specialized method applies at a given level is
+# decided purely by dispatch (on the sentinel type `_NoLogdensityRel`).
+@inline function logdensity_def(μ, ν, x)
+    _logdensity_rel_descent(μ, basemeasure_depth(μ), ν, basemeasure_depth(ν), x)
+end
 
-    for j in 1:(N-1)
-        push!(q.args, :(Δℓ = logdensity_def(νs[$j], x)))
-        # push!(q.args, :(println("Subtracting", Δℓ)))
-        push!(q.args, :(ℓ -= Δℓ))
+@generated function _logdensity_rel_descent(
+    μ,
+    ::StaticInteger{M},
+    ν,
+    ::StaticInteger{N},
+    x,
+) where {M,N}
+    μsym(i) = Symbol(:μ_, i)
+    νsym(j) = Symbol(:ν_, j)
+    prog = Expr(:block, Expr(:meta, :inline), :(μ_0 = μ), :(ν_0 = ν))
+    terms = Any[]
+    n_checks = 0
+    # Return via a specialized `logdensity_rel_def` method for the current
+    # measure pair, if available. Whether one is available is decided purely
+    # by type, so unsuccessful checks are free at run time:
+    function emit_check!(i, j)
+        r = Symbol(:r_, n_checks)
+        n_checks += 1
+        push!(prog.args, :($r = logdensity_rel_def($(μsym(i)), $(νsym(j)), x)))
+        ret = isempty(terms) ? r : :(+($(terms...), $r))
+        push!(prog.args, :(if !($r isa _NoLogdensityRel)
+            return $ret
+        end))
+    end
+    i = j = 0
+    emit_check!(i, j)
+    # Equalize depths, accumulating one-sided density terms:
+    while M - i > N - j
+        ℓ = Symbol(:ℓμ_, i)
+        push!(prog.args, :($ℓ = logdensity_def($(μsym(i)), x)))
+        push!(prog.args, :($(μsym(i + 1)) = basemeasure($(μsym(i)))))
+        push!(terms, ℓ)
+        i += 1
+        emit_check!(i, j)
+    end
+    while N - j > M - i
+        ℓ = Symbol(:ℓν_, j)
+        push!(prog.args, :($ℓ = -logdensity_def($(νsym(j)), x)))
+        push!(prog.args, :($(νsym(j + 1)) = basemeasure($(νsym(j)))))
+        push!(terms, ℓ)
+        j += 1
+        emit_check!(i, j)
+    end
+    # Lockstep descent at equal depth:
+    for _ in 1:(M-i)
+        ℓμ, ℓν = Symbol(:ℓμ_, i), Symbol(:ℓν_, j)
+        push!(prog.args, :($ℓμ = logdensity_def($(μsym(i)), x)))
+        push!(prog.args, :($ℓν = -logdensity_def($(νsym(j)), x)))
+        push!(terms, ℓμ, ℓν)
+        push!(prog.args, :($(μsym(i + 1)) = basemeasure($(μsym(i)))))
+        push!(prog.args, :($(νsym(j + 1)) = basemeasure($(νsym(j)))))
+        i += 1
+        j += 1
+        emit_check!(i, j)
     end
+    # Both measures are at root level now:
+    push!(
+        prog.args,
+        :(r_root = _root_logdensity_rel($(μsym(i)), $(νsym(j)), x)),
+    )
+    ret = isempty(terms) ? :r_root : :(+($(terms...), r_root))
+    push!(prog.args, :(return $ret))
+    return prog
+end
+
+# Root measures of the same type are equal almost everywhere for the
+# purpose of pointwise relative densities:
+_root_logdensity_rel(μ::M, ν::M, x) where {M} = zero(logdensity_def(μ, x))
 
-    push!(q.args, :(return ℓ))
-    return q
+function _root_logdensity_rel(@nospecialize(μ), @nospecialize(ν), @nospecialize(x))
+    throw(
+        ArgumentError(
+            "No method available to compute the log-density between measures with root measures of type $(nameof(typeof(μ))) and $(nameof(typeof(ν)))",
+        ),
+    )
 end
 
 @inline density_rel(μ, ν, x) = exp(logdensity_rel(μ, ν, x))
diff --git a/src/interface.jl b/src/interface.jl
index 6003203d..f7f21004 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -5,7 +5,7 @@ using Reexport
 @reexport using MeasureBase
 
 using MeasureBase: basemeasure_depth, proxy, istrue
-using MeasureBase: insupport, basemeasure_sequence, commonbase
+using MeasureBase: insupport, basemeasure_sequence
 using MeasureBase: transport_to, NoTransport
 
 using DensityInterface: logdensityof
@@ -21,7 +21,6 @@ export basemeasure_depth
 export proxy
 export insupport
 export basemeasure_sequence
-export commonbase
 
 using Test
 
diff --git a/src/primitive.jl b/src/primitive.jl
index f43485c2..0e588334 100644
--- a/src/primitive.jl
+++ b/src/primitive.jl
@@ -11,7 +11,7 @@ measures satisfy the following laws:
 
     logdensity_def(μ::PrimitiveMeasure, x) = 0.0
 
-    logdensity_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0
+    logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0
 """
 abstract type PrimitiveMeasure <: AbstractMeasure end
 
@@ -24,7 +24,7 @@ basemeasure(μ::PrimitiveMeasure) = μ
 
 logdensity_def(::PrimitiveMeasure, x) = static(0.0)
 
-logdensity_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0
+logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0
 
 function Pretty.quoteof(μ::M) where {M<:PrimitiveMeasure}
     :($M())
diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl
index 0e6b171a..7bfdf715 100644
--- a/src/primitives/lebesgue.jl
+++ b/src/primitives/lebesgue.jl
@@ -12,9 +12,9 @@ insupport(::LebesgueBase, x) = true
 
 insupport(::LebesgueBase) = Returns(true)
 
-logdensity_def(::LebesgueBase, ::CountingBase, x) = -Inf
+logdensity_rel_def(::LebesgueBase, ::CountingBase, x) = -Inf
 
-logdensity_def(::CountingBase, ::LebesgueBase, x) = Inf
+logdensity_rel_def(::CountingBase, ::LebesgueBase, x) = Inf
 
 @inline getdof(::LebesgueBase) = static(1)
 
diff --git a/src/schema.jl b/src/schema.jl
deleted file mode 100644
index 70c85577..00000000
--- a/src/schema.jl
+++ /dev/null
@@ -1,34 +0,0 @@
-# Taken from https://github.com/cscherrer/NestedTuples.jl/blob/cd298fd1e5f7e571701a6fee916d2d47c19f32f5/src/typelevel.jl
-
-ntkeys(::Type{NamedTuple{K,V}}) where {K,V} = K
-ntvaltype(::Type{NamedTuple{K,V}}) where {K,V} = V
-
-"""
-    schema(::Type)
-    
-`schema` turns a type into a value that's easier to work with.
-Example:
-    julia> nt = (a=(b=[1,2],c=(d=[3,4],e=[5,6])),f=[7,8]);
-    julia> NT = typeof(nt)
-    NamedTuple{(:a, :f),Tuple{NamedTuple{(:b, :c),Tuple{Array{Int64,1},NamedTuple{(:d, :e),Tuple{Array{Int64,1},Array{Int64,1}}}}},Array{Int64,1}}}
-    julia> schema(NT)
-    (a = (b = Array{Int64,1}, c = (d = Array{Int64,1}, e = Array{Int64,1})), f = Array{Int64,1})
-"""
-function schema end
-
-schema(::NamedTuple{(),Tuple{}}) = NamedTuple()
-schema(::Type{NamedTuple{(),Tuple{}}}) = NamedTuple()
-
-function schema(NT::Type{NamedTuple{names,T}}) where {names,T}
-    return NamedTuple{ntkeys(NT)}(schema(ntvaltype(NT)))
-end
-
-function schema(TT::Type{T}) where {T<:Tuple}
-    return schema.(Tuple(TT.types))
-end
-
-schema(t::T) where {T<:Tuple} = schema(T)
-
-schema(t::T) where {T<:NamedTuple} = schema(T)
-
-schema(T) = T
diff --git a/src/utils.jl b/src/utils.jl
index e169c7c1..06792c67 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -93,36 +93,6 @@ measure of the previous term, and with no repeated entries.
     return filter(!isnothing, Base.Cartesian.@ntuple 10 b)
 end
 
-commonbase(μ, ν) = commonbase(μ, ν, Any)
-
-"""
-    commonbase(μ, ν, T) -> Tuple{StaticInt{i}, StaticInt{j}}
-
-Find minimal (with respect to their sum) `i` and `j` such that there is a method
-
-    logdensity_def(basemeasure_sequence(μ)[i], basemeasure_sequence(ν)[j], ::T)
-
-This is used in `logdensity_rel` to help make that function efficient.
-"""
-@inline function commonbase(μ, ν, ::Type{T}) where {T}
-    return commonbase(basemeasure_sequence(μ), basemeasure_sequence(ν), T)
-end
-
-@generated function commonbase(μ::M, ν::N, ::Type{T}) where {M<:Tuple,N<:Tuple,T}
-    m = schema(M)
-    n = schema(N)
-
-    sols = Iterators.filter(
-        ((i, j),) -> static_hasmethod(logdensity_def, Tuple{m[i],n[j],T}),
-        Iterators.product(1:length(m), 1:length(n)),
-    )
-    isempty(sols) && return :(nothing)
-    minsol = static.(argmin(((i, j),) -> i + j, sols))
-    quote
-        $minsol
-    end
-end
-
 mymap(f, gen::Base.Generator) = mymap(f ∘ gen.f, gen.iter)
 mymap(f, inds...) = Iterators.map(f, inds...)
 
diff --git a/test/test_basics.jl b/test/test_basics.jl
index 11f1a8fe..168bac49 100644
--- a/test/test_basics.jl
+++ b/test/test_basics.jl
@@ -222,6 +222,29 @@ end
     @test logdensity_rel(Lebesgue(), Dirac(0.0) + Lebesgue(), 1.0) == 0.0
 
     @test isnan(logdensity_rel(Dirac(0), Dirac(1), 2))
+
+    # The generic implementation descends the base measure chains of both
+    # measures in lockstep, type-stably and with symbolic cancellation of
+    # shared chain suffixes:
+    let μW = MeasureBase.weightedmeasure(0.7, MeasureBase.StdNormal())
+        StdNormal, StdUniform, StdExponential =
+            MeasureBase.StdNormal, MeasureBase.StdUniform, MeasureBase.StdExponential
+        @test @inferred(logdensity_rel(μW, StdNormal(), 0.5)) ≈ 0.7
+        @test @inferred(logdensity_rel(StdNormal(), μW, 0.5)) ≈ -0.7
+        @test @inferred(logdensity_rel(StdNormal(), StdUniform(), 0.5)) ≈
+              logdensityof(StdNormal(), 0.5)
+        p1 = productmeasure((StdNormal(), StdExponential()))
+        p2 = productmeasure((StdUniform(), StdExponential()))
+        @test @inferred(logdensity_rel(p1, p2, (0.5, 0.5))) ≈
+              logdensityof(StdNormal(), 0.5)
+
+        # Incompatible root measures result in an informative exception:
+        @test_throws ArgumentError logdensity_rel(
+            productmeasure((StdNormal(),)),
+            StdNormal()^1,
+            (0.5,),
+        )
+    end
 end
 
 @testset "Density measures and Radon-Nikodym" begin

From 130d6598ddf156d124fc099396ee95df87ee8af0 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 22:35:17 +0200
Subject: [PATCH 058/122] Make transport from mvstd to vector-marginal products
 type stable

Products over vectors of same-typed unknown-DOF marginals (e.g. binds)
now transport from multivariate standard measures via a typed loop
instead of accumulating into a Vector{Any}; marginal vectors with
abstract element type keep the untyped fallback.

Created by generative AI.
---
 src/combinators/product_transport.jl | 32 +++++++++++++++++++++-------
 test/combinators/bind.jl             |  9 ++++++++
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/src/combinators/product_transport.jl b/src/combinators/product_transport.jl
index 2d78ae2f..da2d3461 100644
--- a/src/combinators/product_transport.jl
+++ b/src/combinators/product_transport.jl
@@ -417,15 +417,31 @@ _marginals_from_mvstd_with_rest_nodof(::Tuple{}, ::StdMeasure, x::AbstractVector
     (), x
 
 function _marginals_from_mvstd_with_rest_nodof(
-    νs::AbstractVector{<:AbstractMeasure},
+    νs::AbstractVector{M},
     μ_inner::StdMeasure,
     x::AbstractVector{<:Real},
-)
-    # ToDo: Check for type stability:
-    ys = Vector{Any}(undef, length(eachindex(νs)))
-    x_rest = x
-    for (i, ν) in zip(eachindex(ys), νs)
-        ys[i], x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x_rest)
+) where {M<:AbstractMeasure}
+    if isconcretetype(M)
+        # Marginals of concrete type produce variates of uniform type, so
+        # the loop below is type stable (the type of the remaining stream
+        # stays invariant under repeated view-taking):
+        idxs = eachindex(νs)
+        y1, x_rest = transport_from_mvstd_with_rest(νs[first(idxs)], μ_inner, x)
+        ys = Vector{typeof(y1)}(undef, length(idxs))
+        ys[begin] = y1
+        j = firstindex(ys) + 1
+        for i in Iterators.drop(idxs, 1)
+            ys[j], x_rest = transport_from_mvstd_with_rest(νs[i], μ_inner, x_rest)
+            j += 1
+        end
+        return ys, x_rest
+    else
+        # Fallback for marginals of mixed type:
+        ys_any = Vector{Any}(undef, length(eachindex(νs)))
+        x_rest = x
+        for (i, ν) in zip(eachindex(ys_any), νs)
+            ys_any[i], x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x_rest)
+        end
+        return [y for y in ys_any], x_rest
     end
-    return [y for y in ys], x_rest
 end
diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl
index f7621037..147ddeb6 100644
--- a/test/combinators/bind.jl
+++ b/test/combinators/bind.jl
@@ -91,6 +91,15 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca
 
         # Variates that are too long must not go unnoticed:
         @test_throws ArgumentError logdensityof(μ, vcat(xy, [0.5]))
+
+        # Products of same-typed unknown-DOF marginals transport type-stably:
+        P = productmeasure([μ, μ])
+        yP = rand(stblrng(), Float64, P)
+        z = transport_to(StdUniform()^6, P)(yP)
+        @test z isa AbstractVector{<:Real} && length(z) == 6
+        yP_reco, rest = MeasureBase.transport_from_mvstd_with_rest(P, StdUniform(), z)
+        @test yP_reco isa Vector{<:AbstractVector{Float64}}
+        @test yP_reco ≈ yP && isempty(rest)
     end
 
     @testset "mbind with merge" begin

From 60284898dc735608a89d27465e49d4073311099e Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 22:45:42 +0200
Subject: [PATCH 059/122] Fix method dispatch ambiguities and test for
 ambiguities

Fixes ambiguities between the static-size power density methods and
powers of primitive measures (the existing disambiguation method did not
cover them), between the massof methods generated by @useproxy and
massof over intervals, and in the legacy kernel constructors. Package
ambiguity testing (including the Aqua ambiguity check) is now enabled in
the test suite.

Created by generative AI.
---
 src/combinators/power.jl              | 11 +++++++++--
 src/combinators/smart-constructors.jl |  4 ++++
 src/proxies.jl                        |  3 +++
 test/test_aqua.jl                     |  7 +++----
 4 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 5145a80c..128771cd 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -182,12 +182,19 @@ massof(m::PowerMeasure) = massof(m.parent)^prod(m.axes)
 
 logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0)
 
-# To avoid ambiguities
+# Disambiguation with the static-size power density methods:
 function logdensity_def(
-    ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0},N}}}, ::Any,
+    ::PowerMeasure{P,Tuple{<:StaticOneToLike{N}}},
+    ::Any,
 ) where {P<:PrimitiveMeasure,N}
     static(0.0)
 end
+function logdensity_def(
+    ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0}}}},
+    ::Any,
+) where {P<:PrimitiveMeasure}
+    static(0.0)
+end
 
 
 @inline mspace_elsize(m::PowerMeasure) = axes2size(m.axes)
diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index e2c4cfe3..e035ca1e 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -324,6 +324,10 @@ end
 
 kernel(::Type{P}, nt::NamedTuple) where {P<:ParameterizedMeasure} = kernel(identity, P, nt)
 
+# Disambiguation:
+kernel(::Type{P}, ::NamedTuple{()}) where {P<:ParameterizedMeasure} =
+    TypedTransitionKernel(constructorof(P), identity)
+
 kernel(::Type{T}; kwargs...) where {T} = kernel(T, NamedTuple(kwargs))
 
 function kernel(::Type{M}, ::NamedTuple{()}) where {M}
diff --git a/src/proxies.jl b/src/proxies.jl
index 109b9973..f2805176 100644
--- a/src/proxies.jl
+++ b/src/proxies.jl
@@ -35,6 +35,9 @@ macro useproxy(M)
 
         @inline $MeasureBase.massof(μ::$M) = massof(proxy(μ))
         @inline $MeasureBase.massof(μ::$M, s) = massof(proxy(μ), s)
+        # Disambiguation with massof(μ, ::AbstractInterval):
+        @inline $MeasureBase.massof(μ::$M, s::$(IntervalSets.AbstractInterval)) =
+            massof(proxy(μ), s)
 
         @inline $MeasureBase.smf(μ::$M, x) = smf(proxy(μ), x)
         @inline $MeasureBase.invsmf(μ::$M, x) = invsmf(proxy(μ), x)
diff --git a/test/test_aqua.jl b/test/test_aqua.jl
index d4546ac2..f683e2c1 100644
--- a/test/test_aqua.jl
+++ b/test/test_aqua.jl
@@ -4,14 +4,13 @@ import Test
 import Aqua
 import MeasureBase
 
-#Test.@testset "Package ambiguities" begin
-#    Test.@test isempty(Test.detect_ambiguities(MeasureBase))
-#end # testset
+Test.@testset "Package ambiguities" begin
+    Test.@test isempty(Test.detect_ambiguities(MeasureBase, recursive = true))
+end # testset
 
 Test.@testset "Aqua tests" begin
     Aqua.test_all(
         MeasureBase,
-        ambiguities = false,
         # Only used by package extensions:
         stale_deps = (ignore = [:ArgCheck, :ArraysOfArrays],),
     )

From b3e4e84513f7ab2b4191a76db371092c76f7cc23 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 22:57:41 +0200
Subject: [PATCH 060/122] Remove the legacy TransitionKernel machinery

Removes kernel.jl (AbstractTransitionKernel and its subtypes, the
kernel/kleisli constructors), the parameterized-measure kernel
constructors and the kernel-based productmeasure methods. MKernel and
mbind supersede this functionality. The basekernel helper stays, it is
independent of the kernel types.

Created by generative AI.
---
 src/MeasureBase.jl                    |   1 -
 src/combinators/product.jl            |  11 +++
 src/combinators/smart-constructors.jl |  85 -------------------
 src/kernel.jl                         | 113 --------------------------
 src/parameterized.jl                  |  21 -----
 test/test_basics.jl                   |   5 --
 6 files changed, 11 insertions(+), 225 deletions(-)
 delete mode 100644 src/kernel.jl

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index e4a3ce42..2dc427b5 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -183,7 +183,6 @@ include("mspace.jl")
 include("getdof.jl")
 include("transport.jl")
 include("proxies.jl")
-include("kernel.jl")
 include("parameterized.jl")
 include("domains.jl")
 include("primitive.jl")
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 9899e0a9..46b4ad3f 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -163,6 +163,17 @@ function _basemeasure(
     productmeasure(mappedarray(basemeasure, mar))
 end
 
+"""
+    MeasureBase.basekernel(f)
+
+For a function `f` that returns a measure, return the function that returns
+the base measure instead, satisfying `basekernel(f)(p) == basemeasure(f(p))`.
+"""
+function basekernel end
+
+basekernel(f) = basemeasure ∘ f
+basekernel(f::Returns) = Returns(basemeasure(f.value))
+
 function _basemeasure(
     μ::ProductMeasure{Base.Generator{I,F}},
     ::Type{B},
diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index e035ca1e..be69d142 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -131,21 +131,6 @@ end
 
 @inline _generic_productmeasure_impl(mar::Base.Generator) = ProductMeasure(mar)
 
-# ToDo: Remove or at least refactor this (ProductMeasure shouldn't take a kernel as its argument).
-
-productmeasure(f, param_maps, pars) = productmeasure(kernel(f, param_maps), pars)
-
-function productmeasure(k::ParameterizedTransitionKernel, pars)
-    productmeasure(k.suff, k.param_maps, pars)
-end
-
-function productmeasure(f::Returns{W}, ::typeof(identity), pars) where {W<:WeightedMeasure}
-    ℓ = _logweight(f.value)
-    base = basemeasure(f.value)
-    newbase = productmeasure(Returns(base), identity, pars)
-    weightedmeasure(length(pars) * ℓ, newbase)
-end
-
 ###############################################################################
 # PushforwardMeasure
 
@@ -280,73 +265,3 @@ function weightedmeasure(ℓ, b::WeightedMeasure)
     weightedmeasure(ℓ + _logweight(b), b.base)
 end
 
-###############################################################################
-# TransitionKernel
-
-# kernel(Normal(μ=2))
-function kernel(μ::M) where {M<:ParameterizedMeasure}
-    kernel(M)
-end
-
-function kernel(d::PowerMeasure)
-    Base.Fix2(powermeasure, d.axes) ∘ kernel(d.parent)
-end
-
-function kernel(f)
-    T = Core.Compiler.return_type(f, Tuple{Any})
-    _kernel(f, T)
-end
-
-function _kernel(f, ::Type{T}) where {T}
-    GenericTransitionKernel(f)
-end
-
-function _kernel(f, ::Type{P}) where {N,P<:ParameterizedMeasure{N}}
-    k = length(N)
-    C = constructorof(P)
-    maps = ntuple(Val(k)) do i
-        x -> @inbounds x[i]
-    end
-
-    kernel(params ∘ f, C, NamedTuple{N}(maps))
-end
-
-kernel(f::F, ::Type{M}; kwargs...) where {F<:Function,M} = kernel(f, M, NamedTuple(kwargs))
-
-function kernel(f::F, ::Type{M}, nt::NamedTuple) where {F<:Function,M}
-    ParameterizedTransitionKernel(M, f, nt)
-end
-
-function kernel(f::F, ::Type{M}, ::NamedTuple{()}) where {F<:Function,M}
-    T = Core.Compiler.return_type(f, Tuple{Any})
-    _kernel(f, M, T)
-end
-
-kernel(::Type{P}, nt::NamedTuple) where {P<:ParameterizedMeasure} = kernel(identity, P, nt)
-
-# Disambiguation:
-kernel(::Type{P}, ::NamedTuple{()}) where {P<:ParameterizedMeasure} =
-    TypedTransitionKernel(constructorof(P), identity)
-
-kernel(::Type{T}; kwargs...) where {T} = kernel(T, NamedTuple(kwargs))
-
-function kernel(::Type{M}, ::NamedTuple{()}) where {M}
-    C = constructorof(M)
-    TypedTransitionKernel(C, identity)
-end
-
-function _kernel(f::F, ::Type{M}, ::Type{NT}) where {M,F,N,NT<:NamedTuple{N}}
-    k = length(N)
-    maps = ntuple(Val(k)) do i
-        x -> @inbounds x[i]
-    end
-
-    ParameterizedTransitionKernel(M, values ∘ f, NamedTuple{N}(maps))
-end
-
-kernel(f::F; kwargs...) where {F<:Function} = kernel(f, NamedTuple(kwargs))
-
-function kernel(f::F, nt::NamedTuple{()}) where {F<:Function}
-    T = Core.Compiler.return_type(f, Tuple{Any})
-    _kernel(f, T)
-end
diff --git a/src/kernel.jl b/src/kernel.jl
deleted file mode 100644
index d6667c7b..00000000
--- a/src/kernel.jl
+++ /dev/null
@@ -1,113 +0,0 @@
-export AbstractTransitionKernel,
-    GenericTransitionKernel, TypedTransitionKernel, ParameterizedTransitionKernel
-
-abstract type AbstractTransitionKernel <: AbstractMeasure end
-
-struct GenericTransitionKernel{F} <: AbstractTransitionKernel
-    f::F
-end
-
-(k::GenericTransitionKernel)(x) = k.f(x)
-
-struct TypedTransitionKernel{M,F} <: AbstractTransitionKernel
-    m::M
-    f::F
-end
-
-(k::TypedTransitionKernel)(x) = (k.m ∘ k.f)(x)
-struct ParameterizedTransitionKernel{M,S,N,T} <: AbstractTransitionKernel
-    m::M
-    suff::S
-    param_maps::NamedTuple{N,T}
-
-    function ParameterizedTransitionKernel(
-        ::Type{M},
-        suff::S,
-        param_maps::NamedTuple{N,T},
-    ) where {M,S,N,T}
-        new{Type{M},S,N,T}(M, suff, param_maps)
-    end
-    function ParameterizedTransitionKernel(
-        m::M,
-        suff::S,
-        param_maps::NamedTuple{N,T},
-    ) where {M,S,N,T}
-        new{M,S,N,T}(m, suff, param_maps)
-    end
-end
-
-"""
-A *kernel* is a function that returns a measure.
-
-    k1 = kernel() do x
-        Normal(x, x^2)
-    end
-
-    k2 = kernel(Normal) do x
-        (μ = x, σ = x^2)
-    end
-
-    k3 = kernel(Normal; μ = identity, σ = abs2)
-
-    k4 = kernel(Normal; μ = first, σ = last) do x
-        (x, x^2)
-    end
-
-    x = randn(); k1(x) == k2(x) == k3(x) == k4(x)
-
-This function is not exported, because "kernel" can have so many other meanings.
-See for example https://github.com/JuliaGaussianProcesses/KernelFunctions.jl for
-another common use of this term.
-
-# Reference
-
-* https://en.wikipedia.org/wiki/Markov_kernel
-"""
-function kernel end
-
-mapcall(t, x) = map(func -> func(x), t)
-
-function (k::ParameterizedTransitionKernel)(x)
-    s = k.suff(x)
-    k.m(; mapcall(k.param_maps, s)...)
-end
-
-(k::AbstractTransitionKernel)(x1, x2, xs...) = k((x1, x2, xs...))
-
-(k::AbstractTransitionKernel)(; kwargs...) = k(NamedTuple(kwargs))
-
-"""
-For any `k::TransitionKernel`, `basekernel` is expected to satisfy
-```
-basekernel(k)(p) == (basemeasure ∘ k)(p)
-```
-
-The main purpose of `basekernel` is to make it efficient to compute
-```
-basemeasure(d::ProductMeasure) == productmeasure(basekernel(d.f), d.xs)
-```
-"""
-function basekernel end
-
-# TODO: Find a way to do better than this
-basekernel(f) = basemeasure ∘ f
-
-basekernel(f::Returns) = Returns(basemeasure(f.value))
-
-function Base.show(io::IO, μ::AbstractTransitionKernel)
-    io = IOContext(io, :compact => true)
-    Pretty.pprint(io, μ)
-end
-
-function Pretty.tile(k::K) where {K<:AbstractTransitionKernel}
-    Pretty.list_layout(
-        Pretty.tile.([getproperty(k, p) for p in propertynames(k)]),
-        prefix = nameof(constructorof(K)),
-    )
-end
-
-const kleisli = kernel
-
-export kleisli
-
-kernel(k::AbstractTransitionKernel) = k
diff --git a/src/parameterized.jl b/src/parameterized.jl
index 8b1c8c88..4412ebdd 100644
--- a/src/parameterized.jl
+++ b/src/parameterized.jl
@@ -24,27 +24,6 @@ function Pretty.tile(d::ParameterizedMeasure{()})
     result
 end
 
-# Allow things like
-#
-# julia> Normal{(:μ,)}(2)
-# Normal(μ = 2,)
-function kernel(::Type{P}) where {N,P<:ParameterizedMeasure{N}}
-    C = constructorof(P)
-    _kernel(C, Val(N))
-end
-
-@inline function _kernel(::Type{C}, ::Val{N}) where {C,N}
-    @inline function f(args::T) where {T<:Tuple}
-        C(NamedTuple{N,T}(args))::C{N,T}
-    end
-
-    @inline function f(arg::T) where {T}
-        C(NamedTuple{N,Tuple{T}}((arg,)))::C{N,Tuple{T}}
-    end
-
-    kernel(f)
-end
-
 function (::Type{P})(nt::NamedTuple{K,T}) where {K,T,N,P<:ParameterizedMeasure{N}}
     C = constructorof(P)
     arg = NamedTuple{N}(nt)
diff --git a/test/test_basics.jl b/test/test_basics.jl
index 168bac49..5db95ad7 100644
--- a/test/test_basics.jl
+++ b/test/test_basics.jl
@@ -75,11 +75,6 @@ testbroken_measures = [
     end
 end
 
-# @testset "TransitionKernel" begin
-#     κ = MeasureBase.kernel(MeasureBase.Dirac, identity)
-#     @test rand(κ(1.1)) == 1.1
-# end
-
 @testset "SpikeMixture" begin
     @test rand(SpikeMixture(Dirac(0), 0.5)) == 0
     @test rand(SpikeMixture(Dirac(1), 1.0)) == 1

From 284c05071a399444055ed350e3893994936af86c Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 23:07:25 +0200
Subject: [PATCH 061/122] Flesh out mspace_elsize and flat-stream variate
 consumption

Adds mspace_elsize methods for standard measures (scalar), Dirac and
powers of scalar-variate measures (static where the axes are static),
and NamedTuple-product variate names based on the marginal names alone.
Flat vector streams now support scalar variates (consuming a single
stream element, also in vcat variate splitting) and multi-rank variates
(consumed in flattened form and reshaped). This enables densities and
transport for binds with scalar-variate primary measures and density
evaluation of multi-dimensional powers inside flat streams.

Created by generative AI.
---
 src/collection_utils.jl     | 16 ++++++++++++++--
 src/combinators/combined.jl |  3 +++
 src/combinators/power.jl    |  7 +++++--
 src/combinators/product.jl  |  2 ++
 src/primitives/dirac.jl     |  2 ++
 src/standard/stdmeasure.jl  |  2 ++
 test/combinators/bind.jl    | 10 ++++++++++
 test/test_basics.jl         | 25 +++++++++++++++++++++++++
 8 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/src/collection_utils.jl b/src/collection_utils.jl
index 6b9de0d5..a4c5e979 100644
--- a/src/collection_utils.jl
+++ b/src/collection_utils.jl
@@ -103,10 +103,22 @@ _cat_measures(a::Tuple, b::AbstractVector) = vcat([a...], b)
 _cat_measures(a::AbstractVector, b::AbstractVector) = vcat(a, b)
 
 
-# Take the beginning of a flat vector stream as a variate of size `sz`:
+# Take the beginning of a flat vector stream as a variate of size `sz`,
+# scalar variates have size `()` and multi-rank variates are reshaped:
 Base.@propagate_inbounds _consume_from_stream(x::AbstractVector, sz::Tuple{IntegerLike}) =
     _split_after(x, sz[1])
 
-function _consume_from_stream(x::AbstractVector, @nospecialize(sz::Tuple))
+Base.@propagate_inbounds function _consume_from_stream(x::AbstractVector, ::Tuple{})
+    idxs = maybestatic_eachindex(x)
+    i_first = maybestatic_first(idxs)
+    x[i_first], _get_or_view(x, i_first + one(i_first), maybestatic_last(idxs))
+end
+
+function _consume_from_stream(x::AbstractVector, sz::Tuple{Vararg{IntegerLike}})
+    a_flat, x_rest = _split_after(x, size2length(sz))
+    return maybestatic_reshape(a_flat, sz), x_rest
+end
+
+function _consume_from_stream(x::AbstractVector, @nospecialize(sz))
     throw(ArgumentError("Can't consume a variate of size $sz from a flat vector stream"))
 end
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index 4325ced8..877e22ee 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -33,6 +33,9 @@ end
 _split_variate_byvalue(::typeof(vcat), test_a::AbstractVector, ab::AbstractVector) =
     _split_after(ab, length(test_a))
 
+_split_variate_byvalue(::typeof(vcat), ::Real, ab::AbstractVector) =
+    _consume_from_stream(ab, ())
+
 _split_variate_byvalue(::typeof(vcat), ::NTuple{N,Any}, ab::Tuple) where {N} =
     _split_after(ab, Val{N}())
 
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 128771cd..0ac4f3b6 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -77,6 +77,11 @@ end
 
 marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes)
 
+# Powers of scalar-variate measures have array-valued variates of known size:
+@inline mspace_elsize(μ::PowerMeasure) = _pwr_mspace_elsize(μ, mspace_elsize(pwr_base(μ)))
+@inline _pwr_mspace_elsize(μ::PowerMeasure, ::Tuple{}) = pwr_size(μ)
+@inline _pwr_mspace_elsize(μ::PowerMeasure, ::Any) = NoMSpaceElementSize{typeof(μ)}()
+
 function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{AbstractArray,N}}) where {N}
     powermeasure(μ, dims)
 end
@@ -196,5 +201,3 @@ function logdensity_def(
     static(0.0)
 end
 
-
-@inline mspace_elsize(m::PowerMeasure) = axes2size(m.axes)
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 46b4ad3f..0e02fdf6 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -89,6 +89,8 @@ proxy(μ::ProductMeasure{<:FillArrays.Fill}) =
     mapreduce(logdensity_rel, +, marginals(μ), marginals(ν), x)
 end
 
+_mspace_names(μ::ProductMeasure{<:NamedTuple{names}}) where {names} = names
+
 function Pretty.tile(d::ProductMeasure{T}) where {T<:Tuple}
     Pretty.list_layout(Pretty.tile.([marginals(d)...]), sep = " ⊗ ")
 end
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index ba044c1d..e4e64a50 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -43,6 +43,8 @@ insupport(d::Dirac, x) = x == d.x
 
 @inline getdof(::Dirac) = static(0)
 
+@inline mspace_elsize(μ::Dirac) = maybestatic_size(μ.x)
+
 @propagate_inbounds function checked_arg(μ::Dirac, x)
     @boundscheck insupport(μ, x) || throw(ArgumentError("Invalid variate for measure"))
     x
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index 81409796..db0e42fe 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -11,6 +11,8 @@ The type of an `N`-dimensional power of a standard measure of type `MU`.
 """
 const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}}
 
+@inline mspace_elsize(::StdMeasure) = ()
+
 @inline check_dof(::StdMeasure, ::StdMeasure) = nothing
 
 @inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x
diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl
index 147ddeb6..1b62e3df 100644
--- a/test/combinators/bind.jl
+++ b/test/combinators/bind.jl
@@ -82,6 +82,16 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca
         @test_throws ArgumentError transport_to(StdUniform()^5, μ)(xy)
         @test_throws ArgumentError transport_to(StdUniform()^2, μ)(xy)
 
+        # Scalar-variate primary measures work in vcat streams:
+        f_βs(a) = pushfwd(Mul(abs(a) + 0.5), StdNormal())^2
+        μsc = mbind(f_βs, StdExponential(), vcat)
+        xys = rand(stblrng(), Float64, μsc)
+        @test xys isa AbstractVector{<:Real} && length(xys) == 3
+        @test logdensityof(μsc, xys) ≈
+              logdensityof(StdExponential(), xys[1]) + logdensityof(f_βs(xys[1]), xys[2:3])
+        ysc = transport_to(StdUniform()^3, μsc)(xys)
+        @test transport_to(μsc, StdUniform()^3)(ysc) ≈ xys
+
         # Nested binds evaluate in a single with-rest pass:
         μnest = mbind(f_βv, μ, vcat)
         xyz = rand(stblrng(), Float64, μnest)
diff --git a/test/test_basics.jl b/test/test_basics.jl
index 5db95ad7..8dbdbe4d 100644
--- a/test/test_basics.jl
+++ b/test/test_basics.jl
@@ -199,6 +199,31 @@ end
     end
 end
 
+@testset "logdensityof_with_rest" begin
+    StdNormal = MeasureBase.StdNormal
+    x = [0.3, 0.7, 0.2, 0.9, 0.5]
+
+    # Scalar variates consume one stream element:
+    @test MeasureBase.mspace_elsize(StdNormal()) == ()
+    ℓ, a, x_rest = MeasureBase.logdensityof_with_rest(StdNormal(), x)
+    @test a == 0.3 && length(x_rest) == 4
+    @test ℓ ≈ logdensityof(StdNormal(), 0.3)
+
+    # Vector variates:
+    @test MeasureBase.mspace_elsize(StdNormal()^2) == (2,)
+    ℓ, a, x_rest = MeasureBase.logdensityof_with_rest(StdNormal()^2, x)
+    @test a == [0.3, 0.7] && length(x_rest) == 3
+    @test ℓ ≈ logdensityof(StdNormal()^2, [0.3, 0.7])
+
+    # Multi-rank variates are consumed in flattened form and reshaped:
+    @test MeasureBase.mspace_elsize(StdNormal()^(2, 2)) == (2, 2)
+    ℓ, a, x_rest = MeasureBase.logdensityof_with_rest(StdNormal()^(2, 2), x)
+    @test a == [0.3 0.2; 0.7 0.9] && length(x_rest) == 1
+    @test ℓ ≈ logdensityof(StdNormal()^(2, 2), a)
+
+    @test MeasureBase.mspace_elsize(Dirac([1, 2])) == (2,)
+end
+
 @testset "logdensity_rel" begin
     @test logdensity_rel(Dirac(0.0) + Lebesgue(), Dirac(1.0), 0.0) == Inf
     @test logdensity_rel(Dirac(0.0) + Lebesgue(), Dirac(1.0), 1.0) == -Inf

From 753def4045619bd343b910e3e836996d09b4586a Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 23:07:25 +0200
Subject: [PATCH 062/122] Move product relative densities to logdensity_rel_def

Relative densities between product measures were a specialization of
logdensity_rel itself, bypassing the support-check layer at the product
level. They are now logdensity_rel_def methods (with a type-stable
tuple-marginals variant) evaluating marginals via unsafe_logdensity_rel,
support checking happens for the products as a whole in logdensity_rel.

Created by generative AI.
---
 src/combinators/product.jl | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 0e02fdf6..22dae865 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -85,8 +85,20 @@ end
 proxy(μ::ProductMeasure{<:FillArrays.Fill}) =
     powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ)))
 
-@inline function logdensity_rel(μ::ProductMeasure, ν::ProductMeasure, x)
-    mapreduce(logdensity_rel, +, marginals(μ), marginals(ν), x)
+# Relative densities between products evaluate marginal-wise. Support
+# checks happen at the logdensity_rel level for the whole products, so the
+# unsafe marginal evaluation suffices here:
+@inline function logdensity_rel_def(μ::ProductMeasure, ν::ProductMeasure, x)
+    mapreduce(unsafe_logdensity_rel, +, marginals(μ), marginals(ν), x)
+end
+
+# For tuples, `mapreduce` has trouble with type inference:
+@inline function logdensity_rel_def(
+    μ::ProductMeasure{<:Tuple},
+    ν::ProductMeasure{<:Tuple},
+    x,
+)
+    sum(map(unsafe_logdensity_rel, marginals(μ), marginals(ν), x))
 end
 
 _mspace_names(μ::ProductMeasure{<:NamedTuple{names}}) where {names} = names

From 81f0488031f85826f01f58e9cf207ee6647cf20f Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 23:07:25 +0200
Subject: [PATCH 063/122] Improve failure behavior of interval massof and
 generic primitive densities

The smf-based interval mass now fails with an informative exception for
measures without a statistical measure function instead of a MethodError
on NoSMF values. The generic non-real-variate log-density of primitive
measures returns a static zero instead of false. test_smf now tolerates
insupport results that are not booleans (NoFastInsupport) and the
logdensity_def docstring points to logdensity_rel_def as the extension
point for specialized relative densities.

Created by generative AI.
---
 src/MeasureBase.jl    | 11 ++++++-----
 src/interface.jl      |  3 ++-
 src/mass-interface.jl | 14 +++++++++++++-
 src/primitive.jl      |  2 +-
 4 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 2dc427b5..c5aee330 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -164,11 +164,12 @@ Compute the log-density of `m1` relative to `m2` at the point `x`, assuming
 `insupport(m1, x)` and `insupport(m2, x)`.
 
 The generic implementation descends the base measure chains of both
-measures in lockstep, so it terminates at the first specialized
-`logdensity_def` method for a pair of base measures (in particular at pairs
-of identical primitive measures) and any shared chain suffix cancels
-symbolically. Measure types may add specialized three-argument methods for
-measure pairs whose relative density can be computed more directly.
+measures in lockstep, so it terminates at the first pair of base measures
+for which a specialized relative density is available (in particular at
+pairs of identical primitive measures) and any shared chain suffix cancels
+symbolically. To provide specialized relative densities for pairs of
+measure types, add methods to [`MeasureBase.logdensity_rel_def`](@ref),
+not to `logdensity_def` itself.
 """
 function logdensity_def end
 
diff --git a/src/interface.jl b/src/interface.jl
index f7f21004..31e68ca0 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -118,7 +118,8 @@ function test_smf(μ, n = 100)
         @assert issorted(p)
         x = invsmf.(μ, p)
         @test issorted(x)
-        @test all(istrue ∘ insupport(μ), x)
+        # insupport may return a non-Bool "don't know" (NoFastInsupport):
+        @test all(x_i -> insupport(μ, x_i) != false, x)
 
         @test all((Finv ∘ F).(x) .≈ x)
 
diff --git a/src/mass-interface.jl b/src/mass-interface.jl
index 59a5db88..33937498 100644
--- a/src/mass-interface.jl
+++ b/src/mass-interface.jl
@@ -119,4 +119,16 @@ in this way, users should add the corresponding `massof` method.
 """
 (m::AbstractMeasure)(s) = massof(m, s)
 
-massof(μ, a_b::AbstractInterval) = smf(μ, rightendpoint(a_b)) - smf(μ, leftendpoint(a_b))
+function massof(μ, a_b::AbstractInterval)
+    _smf_interval_massof(μ, smf(μ, rightendpoint(a_b)), smf(μ, leftendpoint(a_b)))
+end
+
+_smf_interval_massof(μ, smf_r, smf_l) = smf_r - smf_l
+
+function _smf_interval_massof(μ, ::NoSMF, ::NoSMF)
+    throw(
+        ArgumentError(
+            "Can't compute the mass over an interval for a measure of type $(nameof(typeof(μ))), no statistical measure function available",
+        ),
+    )
+end
diff --git a/src/primitive.jl b/src/primitive.jl
index 0e588334..80f24847 100644
--- a/src/primitive.jl
+++ b/src/primitive.jl
@@ -20,7 +20,7 @@ basemeasure(μ::PrimitiveMeasure) = μ
 @inline basemeasure_depth(::PrimitiveMeasure) = static(0)
 
 @inline logdensityof_impl(::PrimitiveMeasure, x::Real) = zero(float(typeof(x)))
-@inline logdensityof_impl(::PrimitiveMeasure, x) = false
+@inline logdensityof_impl(::PrimitiveMeasure, x) = static(0.0)
 
 logdensity_def(::PrimitiveMeasure, x) = static(0.0)
 

From 0dec09c36d737808da4d929a3f1e9830d724b226 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 23:37:29 +0200
Subject: [PATCH 064/122] Remove unused firsttype

Created by generative AI.
---
 .../MeasureBaseDistributionsExt.jl                    |  2 +-
 ext/MeasureBaseForwardDiffExt.jl                      |  4 +---
 src/utils.jl                                          | 11 -----------
 3 files changed, 2 insertions(+), 15 deletions(-)

diff --git a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
index 177ec306..a230386a 100644
--- a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
+++ b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
@@ -20,7 +20,7 @@ using MeasureBase: getdof, checked_arg, massof
 using MeasureBase: transport_to, transport_def, transport_origin, from_origin, to_origin
 using MeasureBase: NoTransportOrigin, NoTransport
 using MeasureBase: Reshape
-using MeasureBase: convert_realtype, firsttype, _fwddiff, @_adignore
+using MeasureBase: convert_realtype, _fwddiff, @_adignore
 import MeasureBase:
     _dist_params_numtype, _trafo_cdf_impl, _trafo_quantile_impl, _trafo_quantile_impl_generic
 using MeasureBase: _pushfront, _pushback, _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log
diff --git a/ext/MeasureBaseForwardDiffExt.jl b/ext/MeasureBaseForwardDiffExt.jl
index 20113c7f..0d91f8b8 100644
--- a/ext/MeasureBaseForwardDiffExt.jl
+++ b/ext/MeasureBaseForwardDiffExt.jl
@@ -3,7 +3,7 @@
 module MeasureBaseForwardDiffExt
 
 using MeasureBase
-using MeasureBase: containsnan, firsttype
+using MeasureBase: containsnan
 import ForwardDiff
 
 function MeasureBase.containsnan(x::ForwardDiff.Dual)
@@ -12,7 +12,5 @@ function MeasureBase.containsnan(x::ForwardDiff.Dual)
     return a || b
 end
 
-MeasureBase.firsttype(::Type{T}, ::Type{<:ForwardDiff.Dual{tag,<:Real,N}}) where {T<:Real,tag,N} =
-    ForwardDiff.Dual{tag,T,N}
 
 end # module MeasureBaseForwardDiffExt
diff --git a/src/utils.jl b/src/utils.jl
index 06792c67..6389343f 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -188,17 +188,6 @@ convert_realtype(::Type{T}, x::Union{Tuple,NamedTuple}) where {T<:Real} =
 convert_realtype(::Type{T}, x::AbstractArray) where {T<:Real} =
     map(Base.Fix1(convert_realtype, T), x)
 
-"""
-    MeasureBase.firsttype(::Type{T}, ::Type{U}) where {T<:Real,U<:Real}
-
-Return the first type, but as a dual number type if the second one is dual.
-"""
-function firsttype end
-
-firsttype(::Type{T}, ::Type{U}) where {T<:Real,U<:Real} = T
-
-
-
 # Distributions implementation hooks:
 function _trafo_cdf_impl end
 function _trafo_quantile_impl end

From a9032fae327b37e90badf3def947fd75ad9f4f9d Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 23:37:29 +0200
Subject: [PATCH 065/122] Add regression tests for products of measures with
 value-dependent variate sizes

Densities and standard-measure transport for products of hierarchical
measures (in vector, tuple and named-tuple marginal form) work through
the with-rest machinery; keep it that way.

Created by generative AI.
---
 test/combinators/bind.jl | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl
index 1b62e3df..efd2cabd 100644
--- a/test/combinators/bind.jl
+++ b/test/combinators/bind.jl
@@ -110,6 +110,24 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca
         yP_reco, rest = MeasureBase.transport_from_mvstd_with_rest(P, StdUniform(), z)
         @test yP_reco isa Vector{<:AbstractVector{Float64}}
         @test yP_reco ≈ yP && isempty(rest)
+        @test logdensityof(P, yP) ≈ logdensityof(μ, yP[1]) + logdensityof(μ, yP[2])
+
+        # Products with value-dependent marginal sizes, in all marginal
+        # container flavors:
+        Pt = productmeasure((μ, μ))
+        yt = rand(stblrng(), Float64, Pt)
+        @test logdensityof(Pt, yt) ≈ logdensityof(μ, yt[1]) + logdensityof(μ, yt[2])
+        zt = transport_to(StdUniform()^6, Pt)(yt)
+        yt_reco = transport_to(Pt, StdUniform()^6)(zt)
+        @test all(map(≈, yt_reco, yt))
+
+        Pnt = productmeasure((a = StdNormal(), b = μ))
+        ynt = rand(stblrng(), Float64, Pnt)
+        @test logdensityof(Pnt, ynt) ≈
+              logdensityof(StdNormal(), ynt.a) + logdensityof(μ, ynt.b)
+        znt = transport_to(StdUniform()^4, Pnt)(ynt)
+        ynt_reco = transport_to(Pnt, StdUniform()^4)(znt)
+        @test ynt_reco.a ≈ ynt.a && ynt_reco.b ≈ ynt.b
     end
 
     @testset "mbind with merge" begin

From 9f251d0008a1ae3dab635bb029719513f5490f61 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 6 Jul 2026 23:37:29 +0200
Subject: [PATCH 066/122] Add combinesets for combining measurable sets
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

combinesets(f_c, α, β) combines two sets along the value combination
semantics of mcombine, with specific set representations where possible:
cartesian products concatenate under vcat and merge, one-dimensional
cartesian powers of equal (singleton-typed) base sets concatenate under
vcat, and implicit domains of measures combine into the implicit domain
of the combined measure. CombinedMeasure now provides mdomain.

Also fixes two latent bugs uncovered by the new tests: setcartprod for
NamedTuple sets had an unbound type parameter and never matched, and
membership tests for CartesianPower used reversed argument order in
Base.in.

Created by generative AI.
---
 src/combinators/combined.jl |  2 +
 src/domains.jl              | 77 +++++++++++++++++++++++++++++++++++--
 test/domains.jl             | 65 +++++++++++++++++++++++++++++++
 test/runtests.jl            |  1 +
 4 files changed, 141 insertions(+), 4 deletions(-)
 create mode 100644 test/domains.jl

diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index 877e22ee..bf4e4291 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -127,6 +127,8 @@ end
 # Bypass `checked_arg`, would require splitting ab:
 @inline checked_arg(::CombinedMeasure, ab) = ab
 
+mdomain(μ::CombinedMeasure) = combinesets(μ.f_c, mdomain(μ.α), mdomain(μ.β))
+
 rootmeasure(μ::CombinedMeasure) = mcombine(μ.f_c, rootmeasure(μ.α), rootmeasure(μ.β))
 
 basemeasure(μ::CombinedMeasure) = mcombine(μ.f_c, basemeasure(μ.α), basemeasure(μ.β))
diff --git a/src/domains.jl b/src/domains.jl
index 7458b450..1f3535ef 100644
--- a/src/domains.jl
+++ b/src/domains.jl
@@ -265,7 +265,8 @@ componentsets(s::CartesianProduct) = s._sets
 
 setcartprod(sets::AbstractArray{<:SetLike}) = CartesianProduct(sets)
 setcartprod(sets::Tuple{Vararg{SetLike}}) = CartesianProduct(sets)
-setcartprod(sets::NamedTuple{names,<:Tuple{Vararg{SetLike}}}) = CartesianProduct(sets)
+setcartprod(sets::NamedTuple{names,<:Tuple{Vararg{SetLike}}}) where {names} =
+    CartesianProduct(sets)
 
 @inline Base.in(x::Tuple{}, s::CartesianProduct{Tuple{}}) = true
 @inline Base.in(x::Tuple{Vararg{Any,N}}, s::CartesianProduct{<:Tuple{Vararg{Any,N}}}) where {N} =
@@ -326,7 +327,7 @@ componentsets(s::CartesianPower) = maybestatic_fill(pwr_base(s), pwr_axes(s))
 function Base.in(x::AbstractArray, s::CartesianPower)
     pwr_size(s) == size(x) ||
         throw(ArgumentError("Size of CartesianPower and given point are incompatible."))
-    isempty(x) ? true : all(Base.Fix1(in, pwr_base(s)), x)::Bool
+    isempty(x) ? true : all(Base.Fix2(in, pwr_base(s)), x)::Bool
 end
 
 Base.isempty(s::CartesianPower) = isempty(pwr_base(s)) || size2length(pwr_size(s)) == 0
@@ -347,11 +348,79 @@ end
 
 Represents a combination of two sets.
 
-User code should not create instances of `CombinedMeasure` directly, but should call
-[`combinesets(f_c, α, β)`](@ref) instead.
+User code should not create instances of `CombinedSet` directly, but should
+call [`combinesets(f_c, α, β)`](@ref) instead.
 """
 struct CombinedSet{FC,MA<:SetLike,MB<:SetLike} <: ValueSet
     f_c::FC
     α::MA
     β::MB
 end
+
+function Base.in(@nospecialize(x), ::CombinedSet)
+    throw(ArgumentError("Cannot test if a value lies within a combined set."))
+end
+
+maybe_in(@nospecialize(x), ::CombinedSet) = true
+
+Base.isempty(s::CombinedSet) = isempty(s.α) || isempty(s.β)
+
+"""
+    combinesets(f_c, α, β)
+
+Combine two sets `α` and `β` into the set of all values `f_c(a, b)` with
+`a ∈ α` and `b ∈ β`.
+
+`f_c` must combine values as described in [`mcombine`](@ref). Uses set
+representations more specific than [`MeasureBase.CombinedSet`](@ref) where
+possible.
+"""
+function combinesets end
+export combinesets
+
+@inline combinesets(f_c, α::SetLike, β::SetLike) = _generic_combinesets(f_c, α, β)
+
+# Combining the implicit domains of two measures yields the implicit domain
+# of the combined measure:
+@inline combinesets(f_c, α::ImplicitDomain, β::ImplicitDomain) =
+    ImplicitDomain(mcombine(f_c, α.m, β.m))
+
+@inline _generic_combinesets(::typeof(firstarg), α::SetLike, β::SetLike) = α
+@inline _generic_combinesets(::typeof(secondarg), α::SetLike, β::SetLike) = β
+@inline _generic_combinesets(::typeof(tuple), α::SetLike, β::SetLike) =
+    setcartprod((α, β))
+@inline _generic_combinesets(f_c::typeof(vcat), α::SetLike, β::SetLike) =
+    _combinesets_cat(f_c, α, β)
+@inline _generic_combinesets(f_c::typeof(merge), α::SetLike, β::SetLike) =
+    _combinesets_cat(f_c, α, β)
+@inline _generic_combinesets(f_c, α::SetLike, β::SetLike) = CombinedSet(f_c, α, β)
+
+_combinesets_cat(
+    ::typeof(vcat),
+    α::CartesianProduct{<:AbstractVector},
+    β::CartesianProduct{<:AbstractVector},
+) = setcartprod(vcat(componentsets(α), componentsets(β)))
+
+_combinesets_cat(
+    ::typeof(merge),
+    α::CartesianProduct{<:NamedTuple},
+    β::CartesianProduct{<:NamedTuple},
+) = setcartprod(merge(componentsets(α), componentsets(β)))
+
+# Concatenating one-dimensional powers of equal base sets yields a longer
+# power. Set equality can typically only be established at runtime, so this
+# simplification only happens when it is decidable from the set types alone:
+function _combinesets_cat(
+    ::typeof(vcat),
+    α::CartesianPower{<:Any,<:Tuple{Any}},
+    β::CartesianPower{<:Any,<:Tuple{Any}},
+)
+    if _static_isequal(pwr_base(α), pwr_base(β)) isa True
+        n = size2length(pwr_size(α)) + size2length(pwr_size(β))
+        setcartpower(pwr_base(α), (n,))
+    else
+        CombinedSet(vcat, α, β)
+    end
+end
+
+_combinesets_cat(f_c, α::SetLike, β::SetLike) = CombinedSet(f_c, α, β)
diff --git a/test/domains.jl b/test/domains.jl
new file mode 100644
index 00000000..8d465eb9
--- /dev/null
+++ b/test/domains.jl
@@ -0,0 +1,65 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using MeasureBase
+using MeasureBase: combinesets, setcartprod, setcartpower
+using MeasureBase: CombinedSet, CartesianProduct, CartesianPower, ImplicitDomain
+using MeasureBase: maybe_in, mdomain, mcombine, mbind, pushfwd
+using MeasureBase: StdNormal, StdUniform, StdExponential
+using MeasureBase: ℝ, ℤ
+using OneTwoMany: firstarg, secondarg
+using AffineMaps: Mul
+
+@testset "domains" begin
+    @testset "combinesets" begin
+        @test combinesets(firstarg, ℝ, ℤ) === ℝ
+        @test combinesets(secondarg, ℝ, ℤ) === ℤ
+
+        s_tuple = combinesets(tuple, ℝ, ℤ)
+        @test s_tuple isa CartesianProduct
+        @test (1.5, 2) ∈ s_tuple
+        @test !((1.5, 2.5) ∈ s_tuple)
+
+        pv = setcartprod([ℝ, ℝ])
+        sv = combinesets(vcat, pv, pv)
+        @test sv isa CartesianProduct
+        @test [1.0, 2.0, 3.0, 4.0] ∈ sv
+
+        snt = combinesets(merge, setcartprod((a = ℝ,)), setcartprod((b = ℤ,)))
+        @test snt isa CartesianProduct
+        @test (a = 1.5, b = 2) ∈ snt
+
+        # One-dimensional powers of equal singleton base sets concatenate:
+        spw = combinesets(vcat, setcartpower(ℝ, (2,)), setcartpower(ℝ, (3,)))
+        @test spw isa CartesianPower
+        @test [1.0, 2.0, 3.0, 4.0, 5.0] ∈ spw
+        @test combinesets(vcat, setcartpower(ℝ, (2,)), setcartpower(ℤ, (3,))) isa
+              CombinedSet
+
+        # No specific representation available:
+        sc = combinesets(vcat, ℝ, setcartpower(ℝ, (2,)))
+        @test sc isa CombinedSet
+        @test maybe_in([1.0, 2.0, 3.0], sc)
+        @test !isempty(sc)
+        @test_throws ArgumentError [1.0, 2.0, 3.0] ∈ sc
+
+        # Implicit domains combine into the implicit domain of the
+        # combined measure:
+        sid = combinesets(
+            vcat,
+            ImplicitDomain(StdNormal()^2),
+            ImplicitDomain(StdUniform()^1),
+        )
+        @test sid isa ImplicitDomain
+        @test maybe_in([1.0, 2.0, 3.0], sid)
+    end
+
+    @testset "mdomain of combined measures" begin
+        f_β(a) = pushfwd(Mul(a[1] + 0.5), StdNormal())^2
+        μc = mcombine(vcat, StdNormal()^2, mbind(f_β, StdExponential()^1, vcat))
+        @test μc isa MeasureBase.CombinedMeasure
+        @test mdomain(μc) isa ImplicitDomain
+        @test maybe_in(rand(Float64, μc), mdomain(μc))
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 6168091a..9bf40147 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -17,6 +17,7 @@ include("test_basics.jl")
 include("getdof.jl")
 include("transport.jl")
 include("smf.jl")
+include("domains.jl")
 
 include("test_mooncake.jl")
 

From deeef51d1cfbd3015b72b23700578d31b2fd1d5c Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 01:23:08 +0200
Subject: [PATCH 067/122] Widen ArraysOfArrays compat to 0.6 and 0.7

Created by generative AI.
---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index c0ba40c8..cdd9199c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -59,7 +59,7 @@ MeasureBaseMooncakeExt = "Mooncake"
 
 [compat]
 ArgCheck = "1, 2"
-ArraysOfArrays = "0.6"
+ArraysOfArrays = "0.6, 0.7"
 ChainRulesCore = "1"
 ChangesOfVariables = "0.1.3"
 Compat = "3.35, 4"

From b0ec47cf13ad4d6a57f3b78f44f9e308c8292f73 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 01:28:29 +0200
Subject: [PATCH 068/122] Add logdensities for batched multi-point density
 evaluation

logdensities(mu, X) computes the log-density of mu at each point in X,
preserving the shape of X. Measure types specialize logdensities_impl.

Power measures unwrap into power axes arguments of the internal machinery
(ordered outermost first), so implementation methods never dispatch on
nested PowerMeasure type signatures. Scalar-variate measures evaluate as
a single flat broadcast; power batches with flat variate storage
(ArrayOfSimilarArrays) additionally fuse the per-point reduction into a
single segmented sum, keeping GPU-backed data on-device.

Created by generative AI.
---
 src/MeasureBase.jl     |   4 +-
 src/density-batched.jl | 199 +++++++++++++++++++++++++++++++++++++++++
 test/Project.toml      |   1 +
 test/logdensities.jl   |  79 ++++++++++++++++
 test/runtests.jl       |   1 +
 5 files changed, 283 insertions(+), 1 deletion(-)
 create mode 100644 src/density-batched.jl
 create mode 100644 test/logdensities.jl

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index c5aee330..4b3adc4e 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -61,7 +61,8 @@ import HeterogeneousComputing
 using HeterogeneousComputing: real_numtype
 
 using ArraysOfArrays:
-    VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays, VectorOfSimilarVectors, flatview
+    ArrayOfSimilarArrays, VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays,
+    VectorOfSimilarVectors, flatview
 
 using OneTwoMany: firstarg, secondarg
 
@@ -204,6 +205,7 @@ include("combinators/weighted.jl")
 include("combinators/superpose.jl")
 include("combinators/product.jl")
 include("combinators/power.jl")
+include("density-batched.jl")
 include("combinators/spikemixture.jl")
 include("combinators/likelihood.jl")
 include("combinators/restricted.jl")
diff --git a/src/density-batched.jl b/src/density-batched.jl
new file mode 100644
index 00000000..e827a179
--- /dev/null
+++ b/src/density-batched.jl
@@ -0,0 +1,199 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+export logdensities
+
+"""
+    logdensities(μ::AbstractMeasure, X::AbstractArray)
+
+Compute the log-density of `μ` at each point in `X`.
+
+Returns an array of the shape of `X`, semantically equivalent to
+`logdensityof.(Ref(μ), X)`. The computation may be fused across points,
+though: power measures with flat variate storage (e.g. based on
+`ArraysOfArrays.ArrayOfSimilarArrays`) evaluate as a single flat broadcast
+plus a segmented reduction over the underlying flat data, compatible with
+GPU-backed storage.
+
+Measure types should specialize [`MeasureBase.logdensities_impl`](@ref)
+instead of `logdensities` itself.
+"""
+function logdensities end
+
+@inline function logdensities(μ::AbstractMeasure, X::AbstractArray)
+    _logdensities(logdensityof_impl, μ, X)
+end
+
+"""
+    MeasureBase.logdensities_impl(μ::AbstractMeasure, X::AbstractArray)
+
+Implements [`logdensities(μ, X)`](@ref logdensities) for arrays `X` of
+plain `μ`-variates. Power measures never reach `logdensities_impl`, their
+power structure is processed generically beforehand.
+
+Measure types that support fused multi-point evaluation should specialize
+`logdensities_impl`. Implementations must preserve the shape of `X` and
+must handle points outside the support of `μ` (the result must be `-Inf`
+at such points).
+
+The default implementation broadcasts the log-density over `X` for
+measures with scalar variates and falls back to a `map` over `X`
+otherwise.
+"""
+function logdensities_impl end
+
+function logdensities_impl(μ::AbstractMeasure, X::AbstractArray)
+    _logdensities_generic(logdensityof_impl, μ, X)
+end
+
+# Batched density machinery, parameterized over the point-level density
+# function `f` (`logdensityof_impl` or `logdensity_def`).
+#
+# `_logdensities(f, μ, X, powers...)` treats each element of `X` as a
+# variate of `(μ^pN)^…^p1` for `powers = (p1, …, pN)` (power axes ordered
+# outermost first, i.e. in the order in which they are encountered when
+# descending into a variate) and returns the density sum for each element,
+# preserving the shape of `X`. Power measures are unwrapped into the power
+# axes arguments before any other dispatch happens, so implementation
+# methods only ever dispatch on plain measure types.
+
+@inline function _logdensities(f::F, μ, X::AbstractArray, powers::Vararg{Any,N}) where {F,N}
+    _logdensities_stripped(f, μ, X, powers...)
+end
+
+@inline function _logdensities(
+    f::F,
+    μ::PowerMeasure,
+    X::AbstractArray,
+    powers::Vararg{Any,N},
+) where {F,N}
+    _logdensities(f, pwr_base(μ), X, powers..., pwr_axes(μ))
+end
+
+@inline function _logdensities_stripped(f::F, μ, X::AbstractArray) where {F}
+    _logdensities_impl(f, μ, X)
+end
+
+function _logdensities_stripped(
+    f::F,
+    μ,
+    X::AbstractArray,
+    p1,
+    powers::Vararg{Any,N},
+) where {F,N}
+    map(x -> _powered_ld(f, μ, x, p1, powers...), X)
+end
+
+function _logdensities_stripped(
+    f::F,
+    μ,
+    X::ArrayOfSimilarArrays{<:Real},
+    p1,
+    powers::Vararg{Any,N},
+) where {F,N}
+    _logdensities_fused(f, μ, X, mspace_elsize(μ), p1, powers...)
+end
+
+# Absolute densities go through the `logdensities_impl` extension point:
+@inline function _logdensities_impl(::typeof(logdensityof_impl), μ, X::AbstractArray)
+    logdensities_impl(μ, X)
+end
+
+@inline function _logdensities_impl(f::F, μ, X::AbstractArray) where {F}
+    _logdensities_generic(f, μ, X)
+end
+
+@inline function _logdensities_generic(f::F, μ, X::AbstractArray) where {F}
+    _logdensities_byelsize(f, μ, X, mspace_elsize(μ))
+end
+
+# Scalar variates evaluate as a single flat broadcast:
+@inline function _logdensities_byelsize(
+    f::F,
+    μ,
+    X::AbstractArray{<:Real},
+    ::Tuple{},
+) where {F}
+    broadcast(Base.Fix1(f, μ), X)
+end
+
+@inline function _logdensities_byelsize(f::F, μ, X::AbstractArray, ::Any) where {F}
+    map(Base.Fix1(f, μ), X)
+end
+
+# Scalar-variate measure with flat variate storage: evaluate as a single
+# flat broadcast followed by a segmented reduction over the per-point
+# power structure:
+function _logdensities_fused(
+    f::F,
+    μ,
+    X::ArrayOfSimilarArrays{<:Real,M},
+    ::Tuple{},
+    powers::Vararg{Any,N},
+) where {F,M,N}
+    sz_inner = _flat_powers_size(powers...)
+    if length(sz_inner) == M
+        X_flat = flatview(X)
+        if ntuple(i -> size(X_flat, i), Val(M)) != sz_inner
+            throw(ArgumentError("Size of variates doesn't match size of power measure"))
+        end
+        ld_flat = broadcast(Base.Fix1(f, μ), X_flat)
+        reshape(sum(ld_flat, dims = ntuple(identity, Val(M))), size(X))
+    else
+        map(x -> _powered_ld(f, μ, x, powers...), X)
+    end
+end
+
+function _logdensities_fused(
+    f::F,
+    μ,
+    X::AbstractArray,
+    ::Any,
+    p1,
+    powers::Vararg{Any,N},
+) where {F,N}
+    map(x -> _powered_ld(f, μ, x, p1, powers...), X)
+end
+
+# The flat size of a variate of `(μ^pN)^…^p1` for a scalar-variate `μ`,
+# innermost power axes vary fastest:
+@inline _flat_powers_size() = ()
+@inline function _flat_powers_size(p1, powers::Vararg{Any,N}) where {N}
+    (_flat_powers_size(powers...)..., axes2size(p1)...)
+end
+
+# Scalar counterpart of `_logdensities`: log-density of `(μ^pN)^…^p1` at a
+# single variate `x`.
+@inline _powered_ld(f::F, μ, x) where {F} = f(μ, x)
+
+@inline function _powered_ld(
+    f::F,
+    μ::PowerMeasure,
+    x,
+    p1,
+    powers::Vararg{Any,N},
+) where {F,N}
+    _powered_ld(f, pwr_base(μ), x, p1, powers..., pwr_axes(μ))
+end
+
+@inline function _powered_ld(f::F, μ, x, p1, powers::Vararg{Any,N}) where {F,N}
+    if axes2size(p1) != maybestatic_size(x)
+        throw(ArgumentError("Size of variate doesn't match size of power measure"))
+    end
+    R = _powered_ld_type(f, μ, x, powers...)
+    if isempty(x)
+        zero(R)::R
+    else
+        # Conversion needed since summation can turn static into dynamic values:
+        convert(R, _powered_ld_sum(f, μ, x, powers...))::R
+    end
+end
+
+@inline _powered_ld_sum(f::F, μ, x) where {F} = sum(Base.Fix1(f, μ), x)
+
+@inline function _powered_ld_sum(f::F, μ, x, p1, powers::Vararg{Any,N}) where {F,N}
+    sum(_logdensities(f, μ, x, p1, powers...))
+end
+
+@inline function _powered_ld_type(f::F, ::MU, x, powers::Vararg{Any,N}) where {F,MU,N}
+    Core.Compiler.return_type(_powered_ld, Tuple{F,MU,eltype(x),map(typeof, powers)...})
+end
diff --git a/test/Project.toml b/test/Project.toml
index 32833feb..19645d48 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -13,6 +13,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899"
diff --git a/test/logdensities.jl b/test/logdensities.jl
new file mode 100644
index 00000000..89ff253f
--- /dev/null
+++ b/test/logdensities.jl
@@ -0,0 +1,79 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using MeasureBase
+using MeasureBase: logdensities, StdNormal, StdUniform
+using ArraysOfArrays: VectorOfSimilarVectors, nestedview, flatview
+using IrrationalConstants: log2π
+import JLArrays
+using JLArrays: JLArray
+
+stdnormal_ld(x) = -(x^2 + log2π) / 2
+
+@testset "logdensities" begin
+    @testset "scalar variates" begin
+        X = randn(10)
+        @test @inferred(logdensities(StdNormal(), X)) ≈ stdnormal_ld.(X)
+        Xm = randn(2, 3)
+        @test logdensities(StdNormal(), Xm) ≈ stdnormal_ld.(Xm)
+    end
+
+    @testset "powers with nested variates" begin
+        m3 = StdNormal()^3
+        X = [randn(3) for _ in 1:10]
+        @test @inferred(logdensities(m3, X)) ≈ [sum(stdnormal_ld, x) for x in X]
+        @test only(logdensities(m3, [X[1]])) ≈ logdensityof(m3, X[1])
+
+        m23 = StdNormal()^(2, 3)
+        X23 = [randn(2, 3) for _ in 1:5]
+        @test logdensities(m23, X23) ≈ [sum(stdnormal_ld, x) for x in X23]
+
+        mpp = (StdNormal()^(2, 3))^4
+        Xpp = [[randn(2, 3) for _ in 1:4] for _ in 1:6]
+        @test logdensities(mpp, Xpp) ≈ [sum(x -> sum(stdnormal_ld, x), xs) for xs in Xpp]
+    end
+
+    @testset "powers with flat variate storage" begin
+        m3 = StdNormal()^3
+        X = VectorOfSimilarVectors(randn(3, 10))
+        @test @inferred(logdensities(m3, X)) ≈
+              vec(sum(stdnormal_ld.(flatview(X)), dims = 1))
+
+        # Power structure may be stored flattened out within each point:
+        mpp = (StdNormal()^(2, 3))^4
+        Xpp = nestedview(randn(2, 3, 4, 7), 3)
+        @test logdensities(mpp, Xpp) ≈ [sum(stdnormal_ld, x) for x in Xpp]
+    end
+
+    @testset "non-scalar-variate fallback" begin
+        mprod = productmeasure((StdUniform(), StdNormal()))
+        X = [(rand(), randn()) for _ in 1:5]
+        @test logdensities(mprod, X) ≈ logdensityof.(Ref(mprod), X)
+    end
+
+    @testset "size mismatch" begin
+        @test_throws ArgumentError logdensities(StdNormal()^3, [randn(3), randn(2)])
+        @test_throws ArgumentError logdensities(
+            StdNormal()^3,
+            VectorOfSimilarVectors(randn(2, 5)),
+        )
+    end
+
+    @testset "GPU array semantics" begin
+        JLArrays.allowscalar(false)
+
+        X = JLArray(randn(100))
+        ld = logdensities(StdNormal(), X)
+        @test ld isa JLArray
+        @test Array(ld) ≈ stdnormal_ld.(Array(X))
+
+        Xb = VectorOfSimilarVectors(JLArray(randn(3, 50)))
+        ldb = logdensities(StdNormal()^3, Xb)
+        @test ldb isa JLArray
+        @test Array(ldb) ≈ vec(sum(stdnormal_ld.(Array(flatview(Xb))), dims = 1))
+
+        xj = JLArray(randn(10))
+        @test logdensityof(StdNormal()^10, xj) ≈ logdensityof(StdNormal()^10, Array(xj))
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 9bf40147..b6073d76 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -15,6 +15,7 @@ include("test_standard.jl")
 include("test_basics.jl")
 
 include("getdof.jl")
+include("logdensities.jl")
 include("transport.jl")
 include("smf.jl")
 include("domains.jl")

From d2ccb547a64e01693c096dedf2abd5b762e665f2 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 01:51:36 +0200
Subject: [PATCH 069/122] Base power measure density evaluation on the batched
 machinery

logdensityof_impl and logdensity_def for PowerMeasure now unwrap the power
structure into the power axes arguments of the batched density machinery.
Nested powers with flat variate storage evaluate as a single fused
broadcast and segmented reduction, GPU-compatible.

Replaces the per-level power density methods, including the static-size
specialization and its disambiguation methods, and removes the now unused
infer_logdensity_type.

Created by generative AI.
---
 src/combinators/power.jl | 38 +++++++-------------------------------
 src/utils.jl             |  5 -----
 2 files changed, 7 insertions(+), 36 deletions(-)

diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 0ac4f3b6..62e6ba4b 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -102,31 +102,13 @@ params(d::PowerMeasure) = params(first(marginals(d)))
     basemeasure(d.parent)^d.axes
 end
 
-for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)]
-    @eval @inline function $head(d::PowerMeasure{M}, x) where {M}
-        parent_m = d.parent
-        sz_parent = axes2size(d.axes)
-        sz_x = maybestatic_size(x)
-        if sz_parent != sz_x
-            throw(ArgumentError("Size of variate doesn't match size of power measure"))
-        end
-        R = infer_logdensity_type($func, parent_m, eltype(x))
-        if isempty(x)
-            return zero(R)::R
-        else
-            # Need to convert since sum can turn static into dynamic values:
-            return convert(R, sum(Base.Fix1($func, parent_m), x))::R
-        end
-    end
+# Power structure is unwrapped into the power axes arguments of the batched
+# density machinery (see density-batched.jl), which fuses evaluation over
+# flat variate storage:
 
-    @eval @inline function $head(
-        d::PowerMeasure{<:Any,Tuple{<:StaticOneToLike{N}}},
-        x,
-    ) where {N}
-        parent = d.parent
-        sum(1:N) do j
-            @inbounds $func(parent, x[j])
-        end
+for head in [:logdensityof_impl, :logdensity_def]
+    @eval @inline function $head(d::PowerMeasure, x)
+        _powered_ld($head, pwr_base(d), x, pwr_axes(d))
     end
 
     @eval @inline function $head(
@@ -187,13 +169,7 @@ massof(m::PowerMeasure) = massof(m.parent)^prod(m.axes)
 
 logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0)
 
-# Disambiguation with the static-size power density methods:
-function logdensity_def(
-    ::PowerMeasure{P,Tuple{<:StaticOneToLike{N}}},
-    ::Any,
-) where {P<:PrimitiveMeasure,N}
-    static(0.0)
-end
+# Disambiguation with the static-zero-size power density method:
 function logdensity_def(
     ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0}}}},
     ::Any,
diff --git a/src/utils.jl b/src/utils.jl
index 6389343f..451d5045 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -101,11 +101,6 @@ function infer_zero(f, args...)
     zero(typeintersect(AbstractFloat, inferred_type))
 end
 
-function infer_logdensity_type(f::F, ::M, ::Type{T}) where {F,M,T}
-    inferred_type = Core.Compiler.return_type(f, Tuple{M,T})
-    return inferred_type
-end
-
 @inline function allequal(f, x::AbstractArray)
     val = f(first(x))
     @simd for xj in x

From 6c0cc298a1a84c7d6b45a718712354aa6787d44f Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 02:28:26 +0200
Subject: [PATCH 070/122] Rename logdensities_impl to batched_logdensityof_impl

Batched protocol functions pair with their scalar counterparts by name:
batched_logdensityof_impl implements logdensities, further batched_*
functions (with-rest, transport, rand) will follow the same scheme.

Created by generative AI.
---
 src/density-batched.jl | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/density-batched.jl b/src/density-batched.jl
index e827a179..6a0aac44 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -14,7 +14,7 @@ though: power measures with flat variate storage (e.g. based on
 plus a segmented reduction over the underlying flat data, compatible with
 GPU-backed storage.
 
-Measure types should specialize [`MeasureBase.logdensities_impl`](@ref)
+Measure types should specialize [`MeasureBase.batched_logdensityof_impl`](@ref)
 instead of `logdensities` itself.
 """
 function logdensities end
@@ -24,14 +24,14 @@ function logdensities end
 end
 
 """
-    MeasureBase.logdensities_impl(μ::AbstractMeasure, X::AbstractArray)
+    MeasureBase.batched_logdensityof_impl(μ::AbstractMeasure, X::AbstractArray)
 
 Implements [`logdensities(μ, X)`](@ref logdensities) for arrays `X` of
-plain `μ`-variates. Power measures never reach `logdensities_impl`, their
+plain `μ`-variates. Power measures never reach `batched_logdensityof_impl`, their
 power structure is processed generically beforehand.
 
 Measure types that support fused multi-point evaluation should specialize
-`logdensities_impl`. Implementations must preserve the shape of `X` and
+`batched_logdensityof_impl`. Implementations must preserve the shape of `X` and
 must handle points outside the support of `μ` (the result must be `-Inf`
 at such points).
 
@@ -39,9 +39,9 @@ The default implementation broadcasts the log-density over `X` for
 measures with scalar variates and falls back to a `map` over `X`
 otherwise.
 """
-function logdensities_impl end
+function batched_logdensityof_impl end
 
-function logdensities_impl(μ::AbstractMeasure, X::AbstractArray)
+function batched_logdensityof_impl(μ::AbstractMeasure, X::AbstractArray)
     _logdensities_generic(logdensityof_impl, μ, X)
 end
 
@@ -70,7 +70,7 @@ end
 end
 
 @inline function _logdensities_stripped(f::F, μ, X::AbstractArray) where {F}
-    _logdensities_impl(f, μ, X)
+    _batched_logdensityof_impl(f, μ, X)
 end
 
 function _logdensities_stripped(
@@ -93,12 +93,12 @@ function _logdensities_stripped(
     _logdensities_fused(f, μ, X, mspace_elsize(μ), p1, powers...)
 end
 
-# Absolute densities go through the `logdensities_impl` extension point:
-@inline function _logdensities_impl(::typeof(logdensityof_impl), μ, X::AbstractArray)
-    logdensities_impl(μ, X)
+# Absolute densities go through the `batched_logdensityof_impl` extension point:
+@inline function _batched_logdensityof_impl(::typeof(logdensityof_impl), μ, X::AbstractArray)
+    batched_logdensityof_impl(μ, X)
 end
 
-@inline function _logdensities_impl(f::F, μ, X::AbstractArray) where {F}
+@inline function _batched_logdensityof_impl(f::F, μ, X::AbstractArray) where {F}
     _logdensities_generic(f, μ, X)
 end
 

From b549bf10ef5e391d37eba74105d2a46874ac2428 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 03:11:13 +0200
Subject: [PATCH 071/122] Widen Real argument types to Number for traced-value
 compatibility

Reactant traced scalars subtype Number, not Real. Widens primitive leaf
density kernels, batched density eltype gates, with-rest stream
signatures, log-weight arguments and numtype conversion sources.
Real stays where realness is semantic (domain membership, numtype
request parameters).

Created by generative AI.
---
 src/collection_utils.jl               |  8 ++++----
 src/combinators/combined.jl           |  2 +-
 src/combinators/product_transport.jl  | 12 ++++++------
 src/combinators/transformedmeasure.jl |  2 +-
 src/combinators/weighted.jl           |  2 +-
 src/density-batched.jl                |  6 +++---
 src/primitive.jl                      |  2 +-
 src/primitives/counting.jl            |  2 +-
 src/primitives/dirac.jl               |  4 ++--
 src/primitives/lebesgue.jl            |  4 ++--
 src/utils.jl                          |  8 ++++----
 11 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/src/collection_utils.jl b/src/collection_utils.jl
index a4c5e979..5b4d915a 100644
--- a/src/collection_utils.jl
+++ b/src/collection_utils.jl
@@ -84,12 +84,12 @@ _fill_value(x::FillArrays.Fill) = x.value
 _fill_axes(x::FillArrays.Fill) = x.axes
 
 
-_flatten_to_rv(VV::AbstractVector{<:AbstractVector{<:Real}}) = flatview(VectorOfArrays(VV))
-_flatten_to_rv(VV::AbstractVector{<:StaticVector{N,<:Real}}) where {N} =
+_flatten_to_rv(VV::AbstractVector{<:AbstractVector{<:Number}}) = flatview(VectorOfArrays(VV))
+_flatten_to_rv(VV::AbstractVector{<:StaticVector{N,<:Number}}) where {N} =
     flatview(VectorOfSimilarArrays(VV))
 
-_flatten_to_rv(VV::VectorOfSimilarVectors{<:Real}) = flatview(VV)
-_flatten_to_rv(VV::VectorOfVectors{<:Real}) = flatview(VV)
+_flatten_to_rv(VV::VectorOfSimilarVectors{<:Number}) = flatview(VV)
+_flatten_to_rv(VV::VectorOfVectors{<:Number}) = flatview(VV)
 
 _flatten_to_rv(::Tuple{}) = []
 _flatten_to_rv(tpl::Tuple{Vararg{AbstractVector}}) = vcat(tpl...)
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index bf4e4291..d425ae0a 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -33,7 +33,7 @@ end
 _split_variate_byvalue(::typeof(vcat), test_a::AbstractVector, ab::AbstractVector) =
     _split_after(ab, length(test_a))
 
-_split_variate_byvalue(::typeof(vcat), ::Real, ab::AbstractVector) =
+_split_variate_byvalue(::typeof(vcat), ::Number, ab::AbstractVector) =
     _consume_from_stream(ab, ())
 
 _split_variate_byvalue(::typeof(vcat), ::NTuple{N,Any}, ab::Tuple) where {N} =
diff --git a/src/combinators/product_transport.jl b/src/combinators/product_transport.jl
index da2d3461..a534300b 100644
--- a/src/combinators/product_transport.jl
+++ b/src/combinators/product_transport.jl
@@ -372,7 +372,7 @@ end
 
 function _split_x_by_marginals_with_rest(
     dofs::Union{Tuple,AbstractVector},
-    x::AbstractVector{<:Real},
+    x::AbstractVector{<:Number},
 )
     x_idxs = maybestatic_eachindex(x)
     first_idxs = _dof_access_firstidxs(dofs, maybestatic_first(x_idxs))
@@ -385,7 +385,7 @@ function _marginals_from_mvstd_with_rest(
     νs,
     dofs::_KnownDOFs,
     μ_inner::StdMeasure,
-    x::AbstractVector{<:Real},
+    x::AbstractVector{<:Number},
 )
     xs, x_rest = _split_x_by_marginals_with_rest(dofs, x)
     μs = map(n -> μ_inner^n, dofs)
@@ -397,7 +397,7 @@ function _marginals_from_mvstd_with_rest(
     νs,
     dofs,
     μ_inner::StdMeasure,
-    x::AbstractVector{<:Real},
+    x::AbstractVector{<:Number},
 )
     _marginals_from_mvstd_with_rest_nodof(νs, μ_inner, x)
 end
@@ -405,7 +405,7 @@ end
 function _marginals_from_mvstd_with_rest_nodof(
     νs::Tuple{Vararg{AbstractMeasure}},
     μ_inner::StdMeasure,
-    x::AbstractVector{<:Real},
+    x::AbstractVector{<:Number},
 )
     # ToDo: Check for type stability, may need a generated function:
     y1, x_rest = transport_from_mvstd_with_rest(νs[1], μ_inner, x)
@@ -413,13 +413,13 @@ function _marginals_from_mvstd_with_rest_nodof(
     return (y1, y2_end...), x_final_rest
 end
 
-_marginals_from_mvstd_with_rest_nodof(::Tuple{}, ::StdMeasure, x::AbstractVector{<:Real}) =
+_marginals_from_mvstd_with_rest_nodof(::Tuple{}, ::StdMeasure, x::AbstractVector{<:Number}) =
     (), x
 
 function _marginals_from_mvstd_with_rest_nodof(
     νs::AbstractVector{M},
     μ_inner::StdMeasure,
-    x::AbstractVector{<:Real},
+    x::AbstractVector{<:Number},
 ) where {M<:AbstractMeasure}
     if isconcretetype(M)
         # Marginals of concrete type produce variates of uniform type, so
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index ee960f9a..8c8ff9e9 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -118,7 +118,7 @@ end
 # end
 
 # TODO: Would profit from custom pullback:
-function _combine_logd_with_ladj(logd_orig::Real, ladj::Real)
+function _combine_logd_with_ladj(logd_orig::Number, ladj::Number)
     logd_result = logd_orig + ladj
     R = typeof(logd_result)
 
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index 124662b6..e70b8e86 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -44,7 +44,7 @@ function Base.:*(k::T, m::AbstractMeasure) where {T<:Number}
     return weightedmeasure(logk, m)
 end
 
-Base.:*(m::AbstractMeasure, k::Real) = k * m
+Base.:*(m::AbstractMeasure, k::Number) = k * m
 
 gentype(μ::WeightedMeasure) = gentype(μ.base)
 
diff --git a/src/density-batched.jl b/src/density-batched.jl
index 6a0aac44..a29effd4 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -86,7 +86,7 @@ end
 function _logdensities_stripped(
     f::F,
     μ,
-    X::ArrayOfSimilarArrays{<:Real},
+    X::ArrayOfSimilarArrays{<:Number},
     p1,
     powers::Vararg{Any,N},
 ) where {F,N}
@@ -110,7 +110,7 @@ end
 @inline function _logdensities_byelsize(
     f::F,
     μ,
-    X::AbstractArray{<:Real},
+    X::AbstractArray{<:Number},
     ::Tuple{},
 ) where {F}
     broadcast(Base.Fix1(f, μ), X)
@@ -126,7 +126,7 @@ end
 function _logdensities_fused(
     f::F,
     μ,
-    X::ArrayOfSimilarArrays{<:Real,M},
+    X::ArrayOfSimilarArrays{<:Number,M},
     ::Tuple{},
     powers::Vararg{Any,N},
 ) where {F,M,N}
diff --git a/src/primitive.jl b/src/primitive.jl
index 80f24847..0a2aba37 100644
--- a/src/primitive.jl
+++ b/src/primitive.jl
@@ -19,7 +19,7 @@ basemeasure(μ::PrimitiveMeasure) = μ
 
 @inline basemeasure_depth(::PrimitiveMeasure) = static(0)
 
-@inline logdensityof_impl(::PrimitiveMeasure, x::Real) = zero(float(typeof(x)))
+@inline logdensityof_impl(::PrimitiveMeasure, x::Number) = zero(float(typeof(x)))
 @inline logdensityof_impl(::PrimitiveMeasure, x) = static(0.0)
 
 logdensity_def(::PrimitiveMeasure, x) = static(0.0)
diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl
index 5a7a998e..776554c5 100644
--- a/src/primitives/counting.jl
+++ b/src/primitives/counting.jl
@@ -12,7 +12,7 @@ struct Counting{T} <: AbstractMeasure
     Counting(supp) = new{Core.Typeof(supp)}(supp)
 end
 
-@inline function logdensityof_impl(μ::Counting, x::Real)
+@inline function logdensityof_impl(μ::Counting, x::Number)
     R = float(typeof(x))
     insupport(μ, x) ? zero(R) : R(-Inf)
 end
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index e4e64a50..14036dff 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -23,14 +23,14 @@ basemeasure(d::Dirac) = CountingBase()
 
 massof(::Dirac) = static(1.0)
 
-function logdensityof_impl(μ::Dirac, x::Real)
+function logdensityof_impl(μ::Dirac, x::Number)
     R = float(typeof(x))
     insupport(μ, x) ? zero(R) : R(-Inf)
 end
 
 logdensityof_impl(μ::Dirac, x) = insupport(μ, x) ? 0.0 : -Inf
 
-logdensity_def(::Dirac, x::Real) = zero(float(typeof(x)))
+logdensity_def(::Dirac, x::Number) = zero(float(typeof(x)))
 logdensity_def(::Dirac, x) = 0.0
 
 Base.rand(::Random.AbstractRNG, T::Type, μ::Dirac) = μ.x
diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl
index 7bfdf715..c66536bd 100644
--- a/src/primitives/lebesgue.jl
+++ b/src/primitives/lebesgue.jl
@@ -18,7 +18,7 @@ logdensity_rel_def(::CountingBase, ::LebesgueBase, x) = Inf
 
 @inline getdof(::LebesgueBase) = static(1)
 
-@inline checked_arg(::LebesgueBase, x::Real) = x
+@inline checked_arg(::LebesgueBase, x::Number) = x
 
 @propagate_inbounds function checked_arg(::LebesgueBase, x::Any)
     @boundscheck throw(ArgumentError("Invalid variate type for measure"))
@@ -63,7 +63,7 @@ insupport(μ::Lebesgue, x) = x ∈ μ.support
 
 insupport(::Lebesgue{RealValues}, ::Real) = true
 
-@inline function logdensityof_impl(μ::Lebesgue, x::Real)
+@inline function logdensityof_impl(μ::Lebesgue, x::Number)
     R = float(typeof(x))
     insupport(μ, x) ? zero(R) : R(-Inf)
 end
diff --git a/src/utils.jl b/src/utils.jl
index 451d5045..b6035e96 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -138,7 +138,7 @@ fcomp(::typeof(identity), g) = g
 fcomp(f, ::typeof(identity)) = f
 fcomp(::typeof(identity), ::typeof(identity)) = identity
 
-near_neg_inf(::Type{T}) where {T<:Real} = T(-1E38) # Still fits into Float32
+near_neg_inf(::Type{T}) where {T<:Number} = T(-1E38) # Still fits into Float32
 
 isneginf(x) = isinf(x) && x < zero(x)
 isposinf(x) = isinf(x) && x > zero(x)
@@ -149,7 +149,7 @@ isapproxzero(A::AbstractArray) = all(isapproxzero, A)
 isapproxone(x::T) where {T<:Real} = x ≈ one(T)
 isapproxone(A::AbstractArray) = all(isapproxone, A)
 
-containsnan(x::Real) = isnan(x)
+containsnan(x::Number) = isnan(x)
 containsnan(x) = any(containsnan, x)
 
 
@@ -176,8 +176,8 @@ function convert_realtype end
 
 @inline convert_realtype(::Type{T}, x::T) where {T<:Real} = x
 @inline convert_realtype(::Type{T}, x::AbstractArray{T}) where {T<:Real} = x
-@inline convert_realtype(::Type{T}, x::U) where {T<:Real,U<:Real} = T(x)
-convert_realtype(::Type{T}, x::AbstractArray{U}) where {T<:Real,U<:Real} = T.(x)
+@inline convert_realtype(::Type{T}, x::U) where {T<:Real,U<:Number} = T(x)
+convert_realtype(::Type{T}, x::AbstractArray{U}) where {T<:Real,U<:Number} = T.(x)
 convert_realtype(::Type{T}, x::Union{Tuple,NamedTuple}) where {T<:Real} =
     map(Base.Fix1(convert_realtype, T), x)
 convert_realtype(::Type{T}, x::AbstractArray) where {T<:Real} =

From 36293fb7d1af19c17f8280e6c6a90d6c8ce569fe Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 12:40:21 +0200
Subject: [PATCH 072/122] Make density kernels and support checks branch-free

Replaces insupport ternaries in the primitive and standard density
kernels with _checksupport (ifelse-based) and the value branches in
_combine_logd_with_ladj with nested ifelse. StdUniform's insupport uses
non-short-circuiting comparisons. Required for traced values (Reactant),
where control flow must not depend on runtime values.

Created by generative AI.
---
 src/combinators/transformedmeasure.jl | 19 ++++++++-----------
 src/primitives/counting.jl            |  4 ++--
 src/primitives/dirac.jl               |  4 ++--
 src/primitives/lebesgue.jl            |  4 ++--
 src/standard/stdexponential.jl        |  4 ++--
 src/standard/stduniform.jl            |  6 +++---
 6 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 8c8ff9e9..706784e6 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -122,17 +122,14 @@ function _combine_logd_with_ladj(logd_orig::Number, ladj::Number)
     logd_result = logd_orig + ladj
     R = typeof(logd_result)
 
-    if isnan(logd_result) && isneginf(logd_orig) && isposinf(ladj)
-        # Zero μ wins against infinite volume:
-        R(-Inf)::R
-    elseif isfinite(logd_orig) && isneginf(ladj)
-        # Maybe  also for isneginf(logd_orig) && isfinite(ladj) ?
-        # Return constant -Inf to prevent problems with ForwardDiff:
-        #R(-Inf)
-        near_neg_inf(R)::R # Avoids AdvancedHMC warnings
-    else
-        logd_result::R
-    end
+    # Zero μ wins against infinite volume:
+    zero_wins = isnan(logd_result) & isneginf(logd_orig) & isposinf(ladj)
+    # Maybe also for isneginf(logd_orig) && isfinite(ladj) ?
+    # Return near_neg_inf instead of constant -Inf to prevent problems
+    # with ForwardDiff and to avoid AdvancedHMC warnings:
+    fades_out = isfinite(logd_orig) & isneginf(ladj)
+
+    ifelse(zero_wins, R(-Inf), ifelse(fades_out, near_neg_inf(R), logd_result))::R
 end
 
 function logdensityof_impl(
diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl
index 776554c5..74e7f8c2 100644
--- a/src/primitives/counting.jl
+++ b/src/primitives/counting.jl
@@ -14,10 +14,10 @@ end
 
 @inline function logdensityof_impl(μ::Counting, x::Number)
     R = float(typeof(x))
-    insupport(μ, x) ? zero(R) : R(-Inf)
+    _checksupport(insupport(μ, x), zero(R))
 end
 
-@inline logdensityof_impl(μ::Counting, x) = insupport(μ, x) ? 0.0 : -Inf
+@inline logdensityof_impl(μ::Counting, x) = _checksupport(insupport(μ, x), 0.0)
 
 @inline logdensity_def(μ::Counting, x) = logdensityof(μ, x)
 
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 14036dff..96077696 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -25,10 +25,10 @@ massof(::Dirac) = static(1.0)
 
 function logdensityof_impl(μ::Dirac, x::Number)
     R = float(typeof(x))
-    insupport(μ, x) ? zero(R) : R(-Inf)
+    _checksupport(insupport(μ, x), zero(R))
 end
 
-logdensityof_impl(μ::Dirac, x) = insupport(μ, x) ? 0.0 : -Inf
+logdensityof_impl(μ::Dirac, x) = _checksupport(insupport(μ, x), 0.0)
 
 logdensity_def(::Dirac, x::Number) = zero(float(typeof(x)))
 logdensity_def(::Dirac, x) = 0.0
diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl
index c66536bd..770ed03e 100644
--- a/src/primitives/lebesgue.jl
+++ b/src/primitives/lebesgue.jl
@@ -65,10 +65,10 @@ insupport(::Lebesgue{RealValues}, ::Real) = true
 
 @inline function logdensityof_impl(μ::Lebesgue, x::Number)
     R = float(typeof(x))
-    insupport(μ, x) ? zero(R) : R(-Inf)
+    _checksupport(insupport(μ, x), zero(R))
 end
 
-@inline logdensityof_impl(μ::Lebesgue, x) = insupport(μ, x) ? 0.0 : -Inf
+@inline logdensityof_impl(μ::Lebesgue, x) = _checksupport(insupport(μ, x), 0.0)
 
 massof(::Lebesgue{RealValues}, s::Interval) = width(s)
 
diff --git a/src/standard/stdexponential.jl b/src/standard/stdexponential.jl
index fda5b53a..1e10bb88 100644
--- a/src/standard/stdexponential.jl
+++ b/src/standard/stdexponential.jl
@@ -4,9 +4,9 @@ export StdExponential
 
 insupport(::StdExponential, x) = x ≥ zero(x)
 
-@inline function logdensityof_impl(::StdExponential, x)
+@inline function logdensityof_impl(d::StdExponential, x)
     R = float(typeof(x))
-    x ≥ zero(R) ? convert(R, -x) : R(-Inf)
+    _checksupport(insupport(d, x), convert(R, -x))
 end
 
 @inline logdensity_def(::StdExponential, x) = -x
diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl
index d0fc236d..f9b07bf1 100644
--- a/src/standard/stduniform.jl
+++ b/src/standard/stduniform.jl
@@ -2,11 +2,11 @@ struct StdUniform <: StdMeasure end
 
 export StdUniform
 
-insupport(::StdUniform, x) = zero(x) ≤ x ≤ one(x)
+insupport(::StdUniform, x) = (zero(x) ≤ x) & (x ≤ one(x))
 
-@inline function logdensityof_impl(::StdUniform, x)
+@inline function logdensityof_impl(d::StdUniform, x)
     R = float(typeof(x))
-    zero(x) ≤ x ≤ one(x) ? zero(R) : R(-Inf)
+    _checksupport(insupport(d, x), zero(R))
 end
 
 @inline logdensity_def(::StdUniform, x) = zero(x)

From 71236f81cb21cc2c4d3ac57d3d0a16b7abe55a37 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 12:40:21 +0200
Subject: [PATCH 073/122] Add direct logdensityof_impl for weighted measures

The weight-shifted density of a support-safe base density is
support-safe, so weighted measures need no explicit support check.
Removes a redundant per-point insupport sweep over the base measure.

Created by generative AI.
---
 src/combinators/weighted.jl | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index e70b8e86..033e334f 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -16,6 +16,12 @@ _logweight(::AbstractMeasure) = 0
     d.logweight
 end
 
+# The weight-shifted density of a support-safe base density is support-safe,
+# no explicit support check required:
+@inline function logdensityof_impl(d::AbstractWeightedMeasure, x)
+    d.logweight + logdensityof_impl(basemeasure(d), x)
+end
+
 function Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractWeightedMeasure) where {T}
     rand(rng, T, basemeasure(μ))
 end

From bf8f7b29d9270f37b35d2cf0ef05899a44446b7a Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 12:40:21 +0200
Subject: [PATCH 074/122] Make Distributions univariate transport traced-value
 compatible

Widens Real argument types to Number in the cdf/quantile/affine
transport machinery and makes the StdUniform transport gateways
branch-free (out-of-support results are masked via ifelse, the quantile
argument is clamped to keep eager evaluation valid). Distributions with
standard-measure or affine transport origins now work with traced
values; families that require inverse incomplete beta/gamma functions
remain host-only.

Created by generative AI.
---
 ext/MeasureBaseDistributionsExt/univariate.jl | 39 ++++++++-----------
 1 file changed, 17 insertions(+), 22 deletions(-)

diff --git a/ext/MeasureBaseDistributionsExt/univariate.jl b/ext/MeasureBaseDistributionsExt/univariate.jl
index 126fe36d..f607eac3 100644
--- a/ext/MeasureBaseDistributionsExt/univariate.jl
+++ b/ext/MeasureBaseDistributionsExt/univariate.jl
@@ -12,21 +12,21 @@
 _dist_params_numtype(d::Distribution) = real_numtype(typeof(Distributions.params(d)))
 
 
-@inline _trafo_cdf(d::Distribution{Univariate,Continuous}, x::Real) =
+@inline _trafo_cdf(d::Distribution{Univariate,Continuous}, x::Number) =
     _trafo_cdf_impl(_dist_params_numtype(d), d, x)
 
-@inline _trafo_cdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Real) =
+@inline _trafo_cdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Number) =
     Distributions.cdf(d, x)
 
 
-@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, u::Real) =
+@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, u::Number) =
     _trafo_quantile_impl(_dist_params_numtype(d), d, u)
 
-@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, u::Real) =
+@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, u::Number) =
     _trafo_quantile_impl_generic(d, u)
 
 
-@inline _trafo_quantile_impl_generic(d::Distribution{Univariate,Continuous}, u::Real) =
+@inline _trafo_quantile_impl_generic(d::Distribution{Univariate,Continuous}, u::Number) =
     Distributions.quantile(d, u)
 
 # Workaround for Beta dist, current quantile implementation only supports Float64:
@@ -50,45 +50,40 @@ end
 end
 
 
-@inline function _result_numtype(d::Distribution{Univariate}, x::T) where {T<:Real}
+@inline function _result_numtype(d::Distribution{Univariate}, x::T) where {T<:Number}
     float(promote_type(T, _dist_params_numtype(d)))
 end
 
 
 @inline function MeasureBase.transport_def(::StdUniform, μ::Distribution{Univariate,Continuous}, x)
     R = _result_numtype(μ, x)
-    if Distributions.insupport(μ, x)
-        y = _trafo_cdf(μ, x)
-        convert(R, y)
-    else
-        convert(R, NaN)
-    end
+    y = _trafo_cdf(μ, x)
+    ifelse(Distributions.insupport(μ, x), convert(R, y), convert(R, NaN))
 end
 
 
 @inline function MeasureBase.transport_def(ν::Distribution{Univariate,Continuous}, ::StdUniform, x::T) where {T}
     R = _result_numtype(ν, x)
     TF = float(T)
-    if 0 <= x <= 1
-        # Avoid x ≈ 0 and x ≈ 1 to avoid infinite variate values for target distributions with infinite support:
-        mod_x = ifelse(x ≈ 0, zero(TF) + eps(TF), ifelse(x ≈ 1, one(TF) - eps(TF), convert(TF, x)))
-        y = _trafo_quantile(ν, mod_x)
-        convert(R, y)
-    else
-        convert(R, NaN)
-    end
+    # Avoid x ≈ 0 and x ≈ 1 to avoid infinite variate values for target
+    # distributions with infinite support, keep the quantile argument valid
+    # for out-of-range x (the result is masked to NaN then):
+    clamped_x = clamp(convert(TF, x), zero(TF), one(TF))
+    mod_x = ifelse(x ≈ 0, zero(TF) + eps(TF), ifelse(x ≈ 1, one(TF) - eps(TF), clamped_x))
+    y = _trafo_quantile(ν, mod_x)
+    ifelse((zero(x) <= x) & (x <= one(x)), convert(R, y), convert(R, NaN))
 end
 
 
 # Use standard measures as transformation origin for scaled/translated equivalents:
 
-function _origin_to_affine(ν::Distribution{Univariate}, y::T) where {T<:Real}
+function _origin_to_affine(ν::Distribution{Univariate}, y::T) where {T<:Number}
     trg_offs, trg_scale = Distributions.location(ν), Distributions.scale(ν)
     x = muladd(y, trg_scale, trg_offs)
     convert(_result_numtype(ν, y), x)
 end
 
-function _affine_to_origin(μ::Distribution{Univariate}, x::T) where {T<:Real}
+function _affine_to_origin(μ::Distribution{Univariate}, x::T) where {T<:Number}
     src_offs, src_scale = Distributions.location(μ), Distributions.scale(μ)
     y = (x - src_offs) / src_scale
     convert(_result_numtype(μ, x), y)

From 6be425ea0d36c6bb06560782c737f76464fc7deb Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 12:40:21 +0200
Subject: [PATCH 075/122] Add MeasureBaseReactantExt

Domain membership methods for traced numbers: Reactant's traced scalars
subtype Number, not Real or Integer, so membership in RealValues and
IntegerValues is decided by their value type parameter.

Created by generative AI.
---
 Project.toml                  |  3 +++
 ext/MeasureBaseReactantExt.jl | 12 ++++++++++++
 2 files changed, 15 insertions(+)
 create mode 100644 ext/MeasureBaseReactantExt.jl

diff --git a/Project.toml b/Project.toml
index cdd9199c..8700a16a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -43,6 +43,7 @@ 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"
 
@@ -56,6 +57,7 @@ MeasureBaseDistributionsMooncakeExt = ["Distributions", "Mooncake"]
 MeasureBaseForwardDiffExt = "ForwardDiff"
 MeasureBaseForwardDiffPullbacksExt = "ForwardDiffPullbacks"
 MeasureBaseMooncakeExt = "Mooncake"
+MeasureBaseReactantExt = "Reactant"
 
 [compat]
 ArgCheck = "1, 2"
@@ -87,6 +89,7 @@ PDMats = "0.11"
 PrettyPrinting = "0.3, 0.4"
 PropertyFunctions = "0.2.2"
 Random = "1"
+Reactant = "0.2"
 Reexport = "1"
 SpecialFunctions = "2.1.4"
 Static = "0.8, 1"
diff --git a/ext/MeasureBaseReactantExt.jl b/ext/MeasureBaseReactantExt.jl
new file mode 100644
index 00000000..53bc6b4b
--- /dev/null
+++ b/ext/MeasureBaseReactantExt.jl
@@ -0,0 +1,12 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseReactantExt
+
+using Reactant: TracedRNumber
+import MeasureBase
+using MeasureBase: RealValues, IntegerValues
+
+Base.in(::TracedRNumber{<:Real}, ::RealValues) = true
+Base.in(::TracedRNumber{<:Integer}, ::IntegerValues) = true
+
+end # module MeasureBaseReactantExt

From e8502833965ad4c31ad7379119ef91c192a522b1 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 17 Sep 2026 21:16:38 +0200
Subject: [PATCH 076/122] Require ArraysOfArrays 1.3

ArraysOfArrays 1.x replaces nestedview with sliced and introduces the
split-mode API that the batched evaluation code will build on.

Created by generative AI.
---
 Project.toml                         | 2 +-
 test/distributions/test_transport.jl | 2 +-
 test/logdensities.jl                 | 4 ++--
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/Project.toml b/Project.toml
index 8700a16a..e65a73b2 100644
--- a/Project.toml
+++ b/Project.toml
@@ -61,7 +61,7 @@ MeasureBaseReactantExt = "Reactant"
 
 [compat]
 ArgCheck = "1, 2"
-ArraysOfArrays = "0.6, 0.7"
+ArraysOfArrays = "1.3"
 ChainRulesCore = "1"
 ChangesOfVariables = "0.1.3"
 Compat = "3.35, 4"
diff --git a/test/distributions/test_transport.jl b/test/distributions/test_transport.jl
index 4ae50831..f4692124 100644
--- a/test/distributions/test_transport.jl
+++ b/test/distributions/test_transport.jl
@@ -32,7 +32,7 @@ include("getjacobian.jl")
     end
 
     reshaped_rand(d::Distribution{Univariate}, n) = rand(d, n)
-    reshaped_rand(d::Distribution{Multivariate}, n) = nestedview(rand(d, n))
+    reshaped_rand(d::Distribution{Multivariate}, n) = sliced(rand(d, n))
 
     function test_dist_trafo_moments(trg, src)
         unshaped(x) = first(torv_and_back(x))
diff --git a/test/logdensities.jl b/test/logdensities.jl
index 89ff253f..2f68f72f 100644
--- a/test/logdensities.jl
+++ b/test/logdensities.jl
@@ -4,7 +4,7 @@ using Test
 
 using MeasureBase
 using MeasureBase: logdensities, StdNormal, StdUniform
-using ArraysOfArrays: VectorOfSimilarVectors, nestedview, flatview
+using ArraysOfArrays: VectorOfSimilarVectors, sliced, flatview
 using IrrationalConstants: log2π
 import JLArrays
 using JLArrays: JLArray
@@ -42,7 +42,7 @@ stdnormal_ld(x) = -(x^2 + log2π) / 2
 
         # Power structure may be stored flattened out within each point:
         mpp = (StdNormal()^(2, 3))^4
-        Xpp = nestedview(randn(2, 3, 4, 7), 3)
+        Xpp = sliced(randn(2, 3, 4, 7), 3)
         @test logdensities(mpp, Xpp) ≈ [sum(stdnormal_ld, x) for x in Xpp]
     end
 

From c773afe7e21dff9904dd0d4f1d973dcf9c6c3551 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 17 Sep 2026 21:16:38 +0200
Subject: [PATCH 077/122] Pin Aqua below 0.8.17 in the test dependencies

Aqua 0.8.17 builds a manifest for the persistent-tasks check by walking
the [deps] of every dependency and fails on packages like IntervalSets
that list a weak dependency (RecipesBase) under [deps] as well.

Created by generative AI.
---
 test/Project.toml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/test/Project.toml b/test/Project.toml
index 19645d48..a33f8f3a 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -29,3 +29,6 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[compat]
+Aqua = "0.8 - 0.8.16"

From b29ce0652a41f5dfac08826f32b4696c6a2f5a3a Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 17 Sep 2026 22:42:39 +0200
Subject: [PATCH 078/122] Extend the variate size contract with mspace_flatsize

mspace_elsize now always reports the shape of the outer array of a
variate, also for powers of measures with array-valued variates. The new
mspace_flatsize reports the size of the flat variate storage, which
composes through nested powers. Both are defined for primitives, Dirac,
weighted, half and restricted measures, reshapes and distributions.

Created by generative AI.
---
 .../distribution_measure.jl                   |  7 ++-
 src/combinators/half.jl                       |  3 ++
 src/combinators/power.jl                      |  6 +--
 src/combinators/reshape.jl                    |  3 ++
 src/combinators/restricted.jl                 |  3 ++
 src/combinators/weighted.jl                   |  3 ++
 src/domains.jl                                |  9 ++++
 src/mspace.jl                                 | 43 +++++++++++++++++-
 src/primitives/counting.jl                    |  6 +++
 src/primitives/dirac.jl                       |  3 +-
 src/primitives/lebesgue.jl                    |  6 +++
 src/standard/stdmeasure.jl                    |  1 +
 test/Project.toml                             |  1 +
 test/distributions/test_distributions.jl      |  1 +
 test/distributions/test_shape_contract.jl     | 14 ++++++
 test/runtests.jl                              |  1 +
 test/shape_contract.jl                        | 45 +++++++++++++++++++
 17 files changed, 147 insertions(+), 8 deletions(-)
 create mode 100644 test/distributions/test_shape_contract.jl
 create mode 100644 test/shape_contract.jl

diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index bcbfd558..19a55f9a 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -59,7 +59,12 @@ end
 
 @inline MeasureBase.massof(::DistributionMeasure) = static(1.0)
 
-@inline MeasureBase.mspace_elsize(m::DistributionMeasure{<:ArrayLikeVariate}) = size(m.obj)
+@inline MeasureBase.mspace_elsize(d::Distribution{Univariate}) = ()
+@inline MeasureBase.mspace_elsize(d::Distribution{<:ArrayLikeVariate}) = size(d)
+@inline MeasureBase.mspace_flatsize(d::Distribution{Univariate}) = ()
+@inline MeasureBase.mspace_flatsize(d::Distribution{<:ArrayLikeVariate}) = size(d)
+@inline MeasureBase.mspace_elsize(m::DistributionMeasure) = MeasureBase.mspace_elsize(m.obj)
+@inline MeasureBase.mspace_flatsize(m::DistributionMeasure) = MeasureBase.mspace_flatsize(m.obj)
 
 @inline MeasureBase.getdof(m::DistributionMeasure{<:ArrayLikeVariate{0}}) = 1
 
diff --git a/src/combinators/half.jl b/src/combinators/half.jl
index c93327b0..c9f86879 100644
--- a/src/combinators/half.jl
+++ b/src/combinators/half.jl
@@ -4,6 +4,9 @@ struct Half{M} <: AbstractMeasure
     parent::M
 end
 
+@inline mspace_elsize(μ::Half) = mspace_elsize(μ.parent)
+@inline mspace_flatsize(μ::Half) = mspace_flatsize(μ.parent)
+
 function Base.show(io::IO, μ::Half)
     print(io, "Half")
     show(io, μ.parent)
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 62e6ba4b..443e4532 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -77,10 +77,8 @@ end
 
 marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes)
 
-# Powers of scalar-variate measures have array-valued variates of known size:
-@inline mspace_elsize(μ::PowerMeasure) = _pwr_mspace_elsize(μ, mspace_elsize(pwr_base(μ)))
-@inline _pwr_mspace_elsize(μ::PowerMeasure, ::Tuple{}) = pwr_size(μ)
-@inline _pwr_mspace_elsize(μ::PowerMeasure, ::Any) = NoMSpaceElementSize{typeof(μ)}()
+@inline mspace_elsize(μ::PowerMeasure) = pwr_size(μ)
+@inline mspace_flatsize(μ::PowerMeasure) = _cat_sizes(mspace_flatsize(pwr_base(μ)), pwr_size(μ))
 
 function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{AbstractArray,N}}) where {N}
     powermeasure(μ, dims)
diff --git a/src/combinators/reshape.jl b/src/combinators/reshape.jl
index f77dbb68..01f5be0d 100644
--- a/src/combinators/reshape.jl
+++ b/src/combinators/reshape.jl
@@ -54,3 +54,6 @@ function mreshape end
 
 mreshape(m::AbstractMeasure, sz::IntegerLike...) = mreshape(m, sz)
 mreshape(m::AbstractMeasure, sz::SizeLike) = pushfwd(Reshape(sz, some_mspace_elsize(m)), m)
+
+@inline mspace_elsize(μ::PushforwardMeasure{<:Reshape}) = μ.f.output_size
+@inline mspace_flatsize(μ::PushforwardMeasure{<:Reshape}) = μ.f.output_size
diff --git a/src/combinators/restricted.jl b/src/combinators/restricted.jl
index e3b66212..063c7f33 100644
--- a/src/combinators/restricted.jl
+++ b/src/combinators/restricted.jl
@@ -3,6 +3,9 @@ struct RestrictedMeasure{P,M} <: AbstractMeasure
     base::M
 end
 
+@inline mspace_elsize(μ::RestrictedMeasure) = mspace_elsize(μ.base)
+@inline mspace_flatsize(μ::RestrictedMeasure) = mspace_flatsize(μ.base)
+
 @inline logdensity_def(d::RestrictedMeasure, x) = logdensity_def(d.base, x)
 
 basemeasure(μ::RestrictedMeasure) = μ.base
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index 033e334f..dfa28e74 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -35,6 +35,9 @@ struct WeightedMeasure{R,M} <: AbstractWeightedMeasure
     base::M
 end
 
+@inline mspace_elsize(μ::WeightedMeasure) = mspace_elsize(μ.base)
+@inline mspace_flatsize(μ::WeightedMeasure) = mspace_flatsize(μ.base)
+
 massof(w::WeightedMeasure) = exp(w.logweight) * massof(w.base)
 
 _logweight(μ::WeightedMeasure) = μ.logweight
diff --git a/src/domains.jl b/src/domains.jl
index 1f3535ef..a6a4e800 100644
--- a/src/domains.jl
+++ b/src/domains.jl
@@ -424,3 +424,12 @@ function _combinesets_cat(
 end
 
 _combinesets_cat(f_c, α::SetLike, β::SetLike) = CombinedSet(f_c, α, β)
+
+
+# Element size of the arrays in array-valued sets, for mspace_elsize:
+@inline _valueset_elsize(::RealValues) = ()
+@inline _valueset_elsize(::IntegerValues) = ()
+@inline _valueset_elsize(::BoundedInts) = ()
+@inline _valueset_elsize(::IntervalSets.AbstractInterval) = ()
+@inline _valueset_elsize(s::CartesianPower) = _cat_sizes(_valueset_elsize(pwr_base(s)), pwr_size(s))
+@inline _valueset_elsize(s) = NoMSpaceElementSize{typeof(s)}()
diff --git a/src/mspace.jl b/src/mspace.jl
index 6227537d..f7fe706c 100644
--- a/src/mspace.jl
+++ b/src/mspace.jl
@@ -12,9 +12,15 @@ struct NoMSpaceElementSize{MU} end
     mspace_elsize(μ)
 
 For a measure `μ` over an array-valued measurable space, return the size of
-the arrays that are the elements of the space.
+the arrays that are the elements of the space, `()` for scalar variates.
 
-May return [`NoMSpaceElementSize{typeof(μ)}()`](@ref).
+The size is static where it is known statically. Returns
+[`NoMSpaceElementSize{typeof(μ)}()`](@ref) if the elements of the space
+are not arrays of one common size, e.g. for structured variates or variates
+whose size depends on the value, or if the size can not be determined
+efficiently.
+
+See also [`MeasureBase.mspace_flatsize`](@ref).
 """
 function mspace_elsize end
 export mspace_elsize
@@ -22,6 +28,31 @@ export mspace_elsize
 @inline mspace_elsize(μ::AbstractMeasure) = NoMSpaceElementSize{typeof(μ)}()
 
 
+"""
+    MeasureBase.mspace_flatsize(μ)
+
+Return the size of the flat storage of a variate of `μ`, `()` for scalar
+variates.
+
+Variates of powers of measures with array-valued variates are nested
+arrays, their flat storage has the size of the inner arrays followed by
+the size of the power. Returns [`NoMSpaceElementSize{typeof(μ)}()`](@ref)
+if the variates of `μ` have no flat storage of a common size.
+
+See also [`mspace_elsize`](@ref).
+"""
+function mspace_flatsize end
+
+@inline mspace_flatsize(μ::AbstractMeasure) = NoMSpaceElementSize{typeof(μ)}()
+
+@inline _cat_sizes(a::SizeLike, b::SizeLike) = canonical_size((_size_dims(a)..., _size_dims(b)...))
+@inline _size_dims(sz::Tuple) = sz
+@inline _size_dims(::StaticArrays.Size{S}) where {S} = map(static, S)
+@inline _cat_sizes(a::NoMSpaceElementSize, ::SizeLike) = a
+@inline _cat_sizes(::SizeLike, b::NoMSpaceElementSize) = b
+@inline _cat_sizes(a::NoMSpaceElementSize, ::NoMSpaceElementSize) = a
+
+
 """
     MeasureBase.some_mspace_elsize(μ::AbstractMeasure)
 
@@ -44,3 +75,11 @@ function some_mspace_elsize end
 @inline _mspace_some_elsize_impl(::AbstractMeasure, sz::SizeLike) = sz
 _mspace_some_elsize_impl(μ::AbstractMeasure, ::NoMSpaceElementSize) =
     maybestatic_size(testvalue(μ))
+
+@inline _value_elsize(::Number) = ()
+@inline _value_elsize(x::AbstractArray) = maybestatic_size(x)
+@inline _value_elsize(x) = NoMSpaceElementSize{typeof(x)}()
+
+@inline _value_flatsize(::Number) = ()
+@inline _value_flatsize(x::AbstractArray{<:Number}) = maybestatic_size(x)
+@inline _value_flatsize(x) = NoMSpaceElementSize{typeof(x)}()
diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl
index 74e7f8c2..0586ba2e 100644
--- a/src/primitives/counting.jl
+++ b/src/primitives/counting.jl
@@ -4,6 +4,9 @@ export Counting, CountingBase
 
 struct CountingBase <: PrimitiveMeasure end
 
+@inline mspace_elsize(::CountingBase) = ()
+@inline mspace_flatsize(::CountingBase) = ()
+
 insupport(::CountingBase, x) = true
 
 struct Counting{T} <: AbstractMeasure
@@ -25,6 +28,9 @@ basemeasure(::Counting) = CountingBase()
 
 Counting() = Counting(ℤ)
 
+@inline mspace_elsize(μ::Counting) = _valueset_elsize(μ.support)
+@inline mspace_flatsize(μ::Counting) = _valueset_elsize(μ.support)
+
 testvalue(::Type{T}, d::Counting) where {T} = testvalue(T, d.support)
 
 proxy(d::Counting) = restrict(in(d.support), CountingBase())
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 96077696..32d5d1de 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -43,7 +43,8 @@ insupport(d::Dirac, x) = x == d.x
 
 @inline getdof(::Dirac) = static(0)
 
-@inline mspace_elsize(μ::Dirac) = maybestatic_size(μ.x)
+@inline mspace_elsize(μ::Dirac) = _value_elsize(μ.x)
+@inline mspace_flatsize(μ::Dirac) = _value_flatsize(μ.x)
 
 @propagate_inbounds function checked_arg(μ::Dirac, x)
     @boundscheck insupport(μ, x) || throw(ArgumentError("Invalid variate for measure"))
diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl
index 770ed03e..5846c765 100644
--- a/src/primitives/lebesgue.jl
+++ b/src/primitives/lebesgue.jl
@@ -4,6 +4,9 @@ export Lebesgue
 
 struct LebesgueBase <: PrimitiveMeasure end
 
+@inline mspace_elsize(::LebesgueBase) = ()
+@inline mspace_flatsize(::LebesgueBase) = ()
+
 massof(::LebesgueBase, s::Interval) = width(s)
 
 testvalue(::LebesgueBase) = 0.0
@@ -48,6 +51,9 @@ gentype(::Lebesgue) = Float64
 
 Lebesgue() = Lebesgue(ℝ)
 
+@inline mspace_elsize(μ::Lebesgue) = _valueset_elsize(μ.support)
+@inline mspace_flatsize(μ::Lebesgue) = _valueset_elsize(μ.support)
+
 testvalue(::Type{T}, d::Lebesgue) where {T} = testvalue(T, d.support)::T
 
 proxy(d::Lebesgue) = restrict(in(d.support), LebesgueBase())
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index db0e42fe..367e5891 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -12,6 +12,7 @@ The type of an `N`-dimensional power of a standard measure of type `MU`.
 const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}}
 
 @inline mspace_elsize(::StdMeasure) = ()
+@inline mspace_flatsize(::StdMeasure) = ()
 
 @inline check_dof(::StdMeasure, ::StdMeasure) = nothing
 
diff --git a/test/Project.toml b/test/Project.toml
index a33f8f3a..81579398 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -11,6 +11,7 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605"
+IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
diff --git a/test/distributions/test_distributions.jl b/test/distributions/test_distributions.jl
index 65fab330..c73d50a7 100644
--- a/test/distributions/test_distributions.jl
+++ b/test/distributions/test_distributions.jl
@@ -14,6 +14,7 @@ using .MeasureBaseDistributionsExt:
 @testset "Distributions extension" begin
     include("test_autodiff_utils.jl")
     include("test_measure_interface.jl")
+    include("test_shape_contract.jl")
     include("test_distribution_measure.jl")
     include("test_standard_dist.jl")
     include("test_standard_uniform.jl")
diff --git a/test/distributions/test_shape_contract.jl b/test/distributions/test_shape_contract.jl
new file mode 100644
index 00000000..dc88dea2
--- /dev/null
+++ b/test/distributions/test_shape_contract.jl
@@ -0,0 +1,14 @@
+using Test
+
+using MeasureBase
+using MeasureBase: mspace_elsize, mspace_flatsize
+using Distributions
+using LinearAlgebra: I
+
+@testset "shape contract for Distributions" begin
+    @test @inferred(mspace_elsize(Normal(1, 2))) === ()
+    @test @inferred(mspace_elsize(asmeasure(Normal(1, 2)))) === ()
+    @test @inferred(mspace_flatsize(asmeasure(MvNormal(zeros(3), I(3))))) == (3,)
+    @test @inferred(mspace_elsize(asmeasure(MvNormal(zeros(3), I(3)))^4)) == (4,)
+    @test @inferred(mspace_flatsize(asmeasure(MvNormal(zeros(3), I(3)))^4)) == (3, 4)
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index b6073d76..114aef51 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -15,6 +15,7 @@ include("test_standard.jl")
 include("test_basics.jl")
 
 include("getdof.jl")
+include("shape_contract.jl")
 include("logdensities.jl")
 include("transport.jl")
 include("smf.jl")
diff --git a/test/shape_contract.jl b/test/shape_contract.jl
new file mode 100644
index 00000000..45d4d96c
--- /dev/null
+++ b/test/shape_contract.jl
@@ -0,0 +1,45 @@
+using Test
+
+using MeasureBase
+using MeasureBase: mspace_elsize, mspace_flatsize, NoMSpaceElementSize
+using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic
+using MeasureBase: Dirac, Lebesgue, Counting, LebesgueBase, CountingBase
+using MeasureBase: mreshape, productmeasure, weightedmeasure, pushfwd, mbind, restrict
+using IntervalSets: (..)
+using StaticArrays: SVector, Size
+using Static: static
+
+@testset "shape contract" begin
+    @testset "mspace_elsize and mspace_flatsize" begin
+        for μ in (StdNormal(), StdUniform(), Lebesgue(), Lebesgue(0..1), Counting(), LebesgueBase(), CountingBase(), Dirac(1.5))
+            @test @inferred(mspace_elsize(μ)) === ()
+            @test @inferred(mspace_flatsize(μ)) === ()
+        end
+
+        @test @inferred(mspace_elsize(StdNormal()^3)) == (3,)
+        @test @inferred(mspace_flatsize(StdNormal()^3)) == (3,)
+        @test @inferred(mspace_elsize(StdNormal()^(2, 3))) == (2, 3)
+        @test @inferred(mspace_flatsize(StdNormal()^(2, 3))) == (2, 3)
+
+        @test @inferred(mspace_elsize((StdNormal()^3)^(2, 4))) == (2, 4)
+        @test @inferred(mspace_flatsize((StdNormal()^3)^(2, 4))) == (3, 2, 4)
+        @test @inferred(mspace_flatsize((StdNormal()^static(3))^static(2))) === Size(3, 2)
+        @test @inferred(mspace_flatsize((StdNormal()^static(3))^2)) === (static(3), 2)
+
+        @test @inferred(mspace_elsize(Dirac([1, 2]))) == (2,)
+        @test @inferred(mspace_flatsize(Dirac([1, 2]))) == (2,)
+        @test @inferred(mspace_flatsize(Dirac(SVector(1, 2)))) === Size(2)
+        @test @inferred(mspace_elsize(Dirac([[1], [2]]))) == (2,)
+        @test @inferred(mspace_flatsize(Dirac([[1], [2]])))  isa NoMSpaceElementSize
+        @test @inferred(mspace_elsize(Dirac((a = 1, b = 2)))) isa NoMSpaceElementSize
+
+        @test @inferred(mspace_elsize(weightedmeasure(0.3, StdNormal()^2))) == (2,)
+        @test @inferred(mspace_flatsize(weightedmeasure(0.3, (StdNormal()^2)^3))) == (2, 3)
+        @test @inferred(mspace_elsize(restrict(x -> x > 0, StdNormal()))) === ()
+        @test @inferred(mspace_elsize(mreshape(StdNormal()^6, (2, 3)))) == (2, 3)
+        @test @inferred(mspace_flatsize(mreshape(StdNormal()^6, (2, 3)))) == (2, 3)
+
+        @test @inferred(mspace_elsize(productmeasure((a = StdNormal(), b = StdUniform())))) isa NoMSpaceElementSize
+        @test @inferred(mspace_flatsize(mbind(x -> StdNormal()^2, StdUniform()))) isa NoMSpaceElementSize
+    end
+end

From a97afd0d9b55a5f7fed83b1335097e372c9cafe6 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 17 Sep 2026 22:42:39 +0200
Subject: [PATCH 079/122] Add preferred_stdmeasure with promotion of standard
 measure types

Every measure family declares the standard measure type its variates are
transported to and from by default. Composite measures combine the
preferences of their components via promote_stdmeasure, which prefers
unbounded standard measures as pivots. Dirac measures accept any standard
measure, measures without a standard-measure transport report
NoStdTransport.

Created by generative AI.
---
 ext/MeasureBaseDistributionsExt/dirichlet.jl  |  2 +
 .../distribution_measure.jl                   |  2 +
 .../standard_dist.jl                          |  1 +
 ext/MeasureBaseDistributionsExt/standardmv.jl |  2 +
 ext/MeasureBaseDistributionsExt/univariate.jl |  7 ++
 src/MeasureBase.jl                            |  1 +
 src/combinators/bind.jl                       |  2 +
 src/combinators/combined.jl                   |  4 +
 src/combinators/half.jl                       |  1 +
 src/standard/stdtraits.jl                     | 99 +++++++++++++++++++
 test/distributions/test_shape_contract.jl     | 22 ++++-
 test/shape_contract.jl                        | 40 ++++++++
 12 files changed, 182 insertions(+), 1 deletion(-)
 create mode 100644 src/standard/stdtraits.jl

diff --git a/ext/MeasureBaseDistributionsExt/dirichlet.jl b/ext/MeasureBaseDistributionsExt/dirichlet.jl
index c60eeecd..98693202 100644
--- a/ext/MeasureBaseDistributionsExt/dirichlet.jl
+++ b/ext/MeasureBaseDistributionsExt/dirichlet.jl
@@ -7,6 +7,8 @@ MeasureBase.getdof(m::DirichletMeasure) = getdof(m.obj)
 
 MeasureBase.transport_origin(d::Dirichlet) = StdUniform()^getdof(d)
 
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Dirichlet}) = StdUniform
+
 
 
 function _dirichlet_beta_trafo(α::Real, β::Real, x::Real)
diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index 19a55f9a..8d3eb8eb 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -66,6 +66,8 @@ end
 @inline MeasureBase.mspace_elsize(m::DistributionMeasure) = MeasureBase.mspace_elsize(m.obj)
 @inline MeasureBase.mspace_flatsize(m::DistributionMeasure) = MeasureBase.mspace_flatsize(m.obj)
 
+@inline MeasureBase.preferred_stdmeasure(::Type{AsMeasure{D}}) where {D<:Distribution} = MeasureBase.preferred_stdmeasure(D)
+
 @inline MeasureBase.getdof(m::DistributionMeasure{<:ArrayLikeVariate{0}}) = 1
 
 # Delegate transport to the wrapped distribution:
diff --git a/ext/MeasureBaseDistributionsExt/standard_dist.jl b/ext/MeasureBaseDistributionsExt/standard_dist.jl
index 010b20a9..e175655c 100644
--- a/ext/MeasureBaseDistributionsExt/standard_dist.jl
+++ b/ext/MeasureBaseDistributionsExt/standard_dist.jl
@@ -42,6 +42,7 @@ for (A, B) in [
     (Normal, StdNormal)
 ]
     @eval begin
+        @inline MeasureBase.preferred_stdmeasure(::Type{<:StandardDist{$A}}) = $B
         @inline MeasureBase.transport_origin(d::StandardDist{$A,0}) = $B()
         @inline MeasureBase.transport_origin(d::StandardDist{$A,N}) where {N} = $B()^size(d)
 
diff --git a/ext/MeasureBaseDistributionsExt/standardmv.jl b/ext/MeasureBaseDistributionsExt/standardmv.jl
index c8e99039..48b9d53e 100644
--- a/ext/MeasureBaseDistributionsExt/standardmv.jl
+++ b/ext/MeasureBaseDistributionsExt/standardmv.jl
@@ -6,6 +6,8 @@ MeasureBase.getdof(m::AsMeasure{<:AbstractMvNormal}) = getdof(m.obj)
 
 MeasureBase.transport_origin(ν::MvNormal) = StandardDist{Normal}(length(ν))
 
+@inline MeasureBase.preferred_stdmeasure(::Type{<:AbstractMvNormal}) = StdNormal
+
 _cholesky_L(A) = cholesky(A).L
 _cholesky_L(A::Diagonal{<:Real}) = Diagonal(sqrt.(diag(A)))
 _cholesky_L(A::PDMats.PDiagMat{<:Real}) = Diagonal(sqrt.(A.diag))
diff --git a/ext/MeasureBaseDistributionsExt/univariate.jl b/ext/MeasureBaseDistributionsExt/univariate.jl
index f607eac3..e34ffaa3 100644
--- a/ext/MeasureBaseDistributionsExt/univariate.jl
+++ b/ext/MeasureBaseDistributionsExt/univariate.jl
@@ -3,6 +3,13 @@
 
 @inline MeasureBase.getdof(::Distribution{Univariate}) = static(1)
 
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Distribution{Univariate,Continuous}}) = StdUniform
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Uniform}) = StdUniform
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Exponential}) = StdExponential
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Logistic}) = StdLogistic
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Normal}) = StdNormal
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Distributions.AffineDistribution{<:Any,<:Any,D}}) where {D} = MeasureBase.preferred_stdmeasure(D)
+
 @inline MeasureBase.check_dof(a::Distribution{Univariate}, b::Distribution{Univariate}) = nothing
 
 
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 4b3adc4e..52d6ec14 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -218,6 +218,7 @@ include("standard/stduniform.jl")
 include("standard/stdexponential.jl")
 include("standard/stdlogistic.jl")
 include("standard/stdnormal.jl")
+include("standard/stdtraits.jl")
 include("combinators/product_transport.jl")
 include("combinators/combined.jl")
 include("combinators/bind.jl")
diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl
index ab7ab708..f1d293c4 100644
--- a/src/combinators/bind.jl
+++ b/src/combinators/bind.jl
@@ -241,6 +241,8 @@ function _bind_tpm_sc(::Type{Pair}, μ::Bind, xy::Pair)
 end
 
 const _BindBy{FC} = Bind{<:Any,<:AbstractMeasure,FC}
+
+@inline preferred_stdmeasure(::Type{<:Bind{<:Any,M}}) where {M} = preferred_stdmeasure(M)
 _bind_tpm_sc(f_c::typeof(vcat), μ::_BindBy{typeof(vcat)}, xy::AbstractVector) =
     _bind_tpm_sc_cat(f_c, μ, xy)
 _bind_tpm_sc(f_c::typeof(merge), μ::_BindBy{typeof(merge)}, xy::NamedTuple) =
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index d425ae0a..3d4d3192 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -118,6 +118,10 @@ struct CombinedMeasure{FC,MA<:AbstractMeasure,MB<:AbstractMeasure} <: AbstractMe
     β::MB
 end
 
+@inline function preferred_stdmeasure(::Type{<:CombinedMeasure{<:Any,MA,MB}}) where {MA,MB}
+    promote_stdmeasure(preferred_stdmeasure(MA), preferred_stdmeasure(MB))
+end
+
 
 @inline insupport(μ::CombinedMeasure, ab) = NoFastInsupport{typeof(μ)}()
 
diff --git a/src/combinators/half.jl b/src/combinators/half.jl
index c9f86879..906978ca 100644
--- a/src/combinators/half.jl
+++ b/src/combinators/half.jl
@@ -6,6 +6,7 @@ end
 
 @inline mspace_elsize(μ::Half) = mspace_elsize(μ.parent)
 @inline mspace_flatsize(μ::Half) = mspace_flatsize(μ.parent)
+@inline preferred_stdmeasure(::Type{<:Half{M}}) where {M} = preferred_stdmeasure(M)
 
 function Base.show(io::IO, μ::Half)
     print(io, "Half")
diff --git a/src/standard/stdtraits.jl b/src/standard/stdtraits.jl
new file mode 100644
index 00000000..69752356
--- /dev/null
+++ b/src/standard/stdtraits.jl
@@ -0,0 +1,99 @@
+"""
+    struct MeasureBase.NoStdTransport{MU}
+
+Indicates that measures of type `MU` can't be transported to or from a
+standard measure.
+"""
+struct NoStdTransport{MU} end
+
+"""
+    struct MeasureBase.AnyStdMeasure
+
+Indicates that any standard measure serves as transport partner, e.g.
+for measures with zero degrees of freedom.
+"""
+struct AnyStdMeasure end
+
+const _StdTransportPartner = Union{Type{<:StdMeasure},Type{AnyStdMeasure},Type{<:NoStdTransport}}
+
+"""
+    MeasureBase.preferred_stdmeasure(μ)::Type
+    MeasureBase.preferred_stdmeasure(::Type{MU})::Type
+
+The type of standard measure that variates of `μ` are transported to and
+from by default.
+
+Returns `MeasureBase.AnyStdMeasure` if any standard measure serves and
+`MeasureBase.NoStdTransport{MU}` if measures of type `MU` have no
+standard-measure transport. Composite measures combine the preferences of
+their components via [`MeasureBase.promote_stdmeasure`](@ref).
+
+Measure types that support transport to and from standard measures should
+specialize the type-based method.
+"""
+function preferred_stdmeasure end
+
+@inline preferred_stdmeasure(μ) = preferred_stdmeasure(typeof(μ))
+@inline preferred_stdmeasure(::Type{MU}) where {MU} = NoStdTransport{MU}
+
+@inline preferred_stdmeasure(::Type{MU}) where {MU<:StdMeasure} = MU
+
+"""
+    MeasureBase.promote_stdmeasure(A::Type, B::Type)::Type
+
+Combine two results of [`MeasureBase.preferred_stdmeasure`](@ref) into
+the preferred standard measure type of a measure composed of both.
+
+Standard measure types promote to the one with the wider range of
+values that remain distinguishable in floating point arithmetic:
+`StdUniform` promotes to any other standard measure type,
+`StdExponential` to `StdLogistic` and `StdNormal`, and `StdLogistic`
+to `StdNormal`.
+"""
+function promote_stdmeasure end
+
+@inline function promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:StdMeasure,B<:StdMeasure}
+    ifelse(_stdmeasure_rank(A) >= _stdmeasure_rank(B), A, B)
+end
+
+@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{B}) where {B} = B
+@inline promote_stdmeasure(::Type{A}, ::Type{AnyStdMeasure}) where {A} = A
+@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{AnyStdMeasure}) = AnyStdMeasure
+@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:NoStdTransport,B} = A
+@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A,B<:NoStdTransport} = B
+@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:NoStdTransport,B<:NoStdTransport} = A
+@inline promote_stdmeasure(::Type{A}, ::Type{AnyStdMeasure}) where {A<:NoStdTransport} = A
+@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{B}) where {B<:NoStdTransport} = B
+
+@inline promote_stdmeasure(::Type{A}) where {A} = A
+@inline function promote_stdmeasure(::Type{A}, ::Type{B}, Cs::Vararg{Type,N}) where {A,B,N}
+    promote_stdmeasure(promote_stdmeasure(A, B), Cs...)
+end
+
+@inline _stdmeasure_rank(::Type{StdUniform}) = 1
+@inline _stdmeasure_rank(::Type{StdExponential}) = 2
+@inline _stdmeasure_rank(::Type{StdLogistic}) = 3
+@inline _stdmeasure_rank(::Type{StdNormal}) = 4
+
+
+@inline preferred_stdmeasure(::Type{<:PowerMeasure{M}}) where {M} = preferred_stdmeasure(M)
+@inline preferred_stdmeasure(::Type{<:WeightedMeasure{<:Any,M}}) where {M} = preferred_stdmeasure(M)
+@inline preferred_stdmeasure(::Type{<:RestrictedMeasure{<:Any,M}}) where {M} = preferred_stdmeasure(M)
+@inline preferred_stdmeasure(::Type{<:PushforwardMeasure{<:Any,<:Any,M}}) where {M} = preferred_stdmeasure(M)
+@inline preferred_stdmeasure(::Type{<:Dirac}) = AnyStdMeasure
+
+@inline preferred_stdmeasure(::Type{<:ProductMeasure{M}}) where {M<:AbstractArray} = preferred_stdmeasure(eltype(M))
+
+@inline function preferred_stdmeasure(::Type{<:ProductMeasure{M}}) where {M<:Tuple}
+    _promote_stdmeasure_oftypes(M)
+end
+
+@inline function preferred_stdmeasure(::Type{<:ProductMeasure{NamedTuple{names,M}}}) where {names,M<:Tuple}
+    _promote_stdmeasure_oftypes(M)
+end
+
+@inline _promote_stdmeasure_oftypes(::Type{Tuple{}}) = AnyStdMeasure
+@generated function _promote_stdmeasure_oftypes(::Type{M}) where {M<:Tuple}
+    args = [:(preferred_stdmeasure($T)) for T in M.parameters]
+    :(promote_stdmeasure($(args...)))
+end
diff --git a/test/distributions/test_shape_contract.jl b/test/distributions/test_shape_contract.jl
index dc88dea2..d02c5680 100644
--- a/test/distributions/test_shape_contract.jl
+++ b/test/distributions/test_shape_contract.jl
@@ -1,7 +1,8 @@
 using Test
 
 using MeasureBase
-using MeasureBase: mspace_elsize, mspace_flatsize
+using MeasureBase: mspace_elsize, mspace_flatsize, preferred_stdmeasure, NoStdTransport
+using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic, Dirac, productmeasure
 using Distributions
 using LinearAlgebra: I
 
@@ -11,4 +12,23 @@ using LinearAlgebra: I
     @test @inferred(mspace_flatsize(asmeasure(MvNormal(zeros(3), I(3))))) == (3,)
     @test @inferred(mspace_elsize(asmeasure(MvNormal(zeros(3), I(3)))^4)) == (4,)
     @test @inferred(mspace_flatsize(asmeasure(MvNormal(zeros(3), I(3)))^4)) == (3, 4)
+
+    @test @inferred(preferred_stdmeasure(Normal(1, 2))) === StdNormal
+    @test @inferred(preferred_stdmeasure(asmeasure(Normal(1, 2)))) === StdNormal
+    @test @inferred(preferred_stdmeasure(3 + 2 * Normal())) === StdNormal
+    @test @inferred(preferred_stdmeasure(Uniform(1, 2))) === StdUniform
+    @test @inferred(preferred_stdmeasure(Exponential(2.0))) === StdExponential
+    @test @inferred(preferred_stdmeasure(Logistic(1, 2))) === StdLogistic
+    @test @inferred(preferred_stdmeasure(Beta(2, 3))) === StdUniform
+    @test @inferred(preferred_stdmeasure(truncated(Normal(), 0, 1))) === StdUniform
+    @test @inferred(preferred_stdmeasure(MvNormal(zeros(2), I(2)))) === StdNormal
+    @test @inferred(preferred_stdmeasure(Dirichlet([1.0, 2.0]))) === StdUniform
+    @test @inferred(preferred_stdmeasure(Poisson(3))) <: NoStdTransport
+    @test @inferred(preferred_stdmeasure(StandardDist{Normal}(3))) === StdNormal
+    @test @inferred(preferred_stdmeasure(StandardDist{Uniform}())) === StdUniform
+
+    @test @inferred(preferred_stdmeasure(productmeasure((asmeasure(Beta(2, 3)), asmeasure(Normal()))))) === StdNormal
+    @test @inferred(preferred_stdmeasure(productmeasure((a = Dirac(1.0), b = asmeasure(Beta(2, 3)))))) === StdUniform
+    @test @inferred(preferred_stdmeasure(productmeasure((a = asmeasure(Poisson(2)), b = asmeasure(Beta(2, 3)))))) <: NoStdTransport
+    @test @inferred(preferred_stdmeasure(productmeasure([asmeasure(Normal(i, 1)) for i in 1:3]))) === StdNormal
 end
diff --git a/test/shape_contract.jl b/test/shape_contract.jl
index 45d4d96c..31920c9f 100644
--- a/test/shape_contract.jl
+++ b/test/shape_contract.jl
@@ -2,6 +2,7 @@ using Test
 
 using MeasureBase
 using MeasureBase: mspace_elsize, mspace_flatsize, NoMSpaceElementSize
+using MeasureBase: preferred_stdmeasure, promote_stdmeasure, AnyStdMeasure, NoStdTransport
 using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic
 using MeasureBase: Dirac, Lebesgue, Counting, LebesgueBase, CountingBase
 using MeasureBase: mreshape, productmeasure, weightedmeasure, pushfwd, mbind, restrict
@@ -42,4 +43,43 @@ using Static: static
         @test @inferred(mspace_elsize(productmeasure((a = StdNormal(), b = StdUniform())))) isa NoMSpaceElementSize
         @test @inferred(mspace_flatsize(mbind(x -> StdNormal()^2, StdUniform()))) isa NoMSpaceElementSize
     end
+
+    @testset "preferred_stdmeasure" begin
+        for S in (StdNormal, StdUniform, StdExponential, StdLogistic)
+            @test @inferred(preferred_stdmeasure(S())) === S
+            @test @inferred(preferred_stdmeasure(S()^3)) === S
+            @test @inferred(preferred_stdmeasure(weightedmeasure(0.1, S()))) === S
+            @test @inferred(preferred_stdmeasure(pushfwd(exp, S()))) === S
+            @test @inferred(preferred_stdmeasure(restrict(x -> x > 0, S()))) === S
+        end
+
+        @test @inferred(preferred_stdmeasure(Dirac(2.0))) === AnyStdMeasure
+        @test @inferred(preferred_stdmeasure(Lebesgue())) <: NoStdTransport
+        @test @inferred(preferred_stdmeasure(Counting())) <: NoStdTransport
+
+        @test @inferred(preferred_stdmeasure(productmeasure((StdUniform(), StdNormal())))) === StdNormal
+        @test @inferred(preferred_stdmeasure(productmeasure((a = StdUniform(), b = StdExponential())))) === StdExponential
+        @test @inferred(preferred_stdmeasure(productmeasure((a = Dirac(1.0), b = StdUniform())))) === StdUniform
+        @test @inferred(preferred_stdmeasure(productmeasure((a = Lebesgue(), b = StdUniform())))) <: NoStdTransport
+        @test @inferred(preferred_stdmeasure(productmeasure(fill(StdLogistic(), 3)))) === StdLogistic
+        @test @inferred(preferred_stdmeasure(productmeasure(()))) === AnyStdMeasure
+        @test @inferred(preferred_stdmeasure(mbind(x -> StdNormal()^2, StdUniform()))) === StdUniform
+    end
+
+    @testset "promote_stdmeasure" begin
+        @test @inferred(promote_stdmeasure(StdUniform, StdNormal)) === StdNormal
+        @test @inferred(promote_stdmeasure(StdNormal, StdUniform)) === StdNormal
+        @test @inferred(promote_stdmeasure(StdUniform, StdExponential)) === StdExponential
+        @test @inferred(promote_stdmeasure(StdExponential, StdLogistic)) === StdLogistic
+        @test @inferred(promote_stdmeasure(StdLogistic, StdNormal)) === StdNormal
+        @test @inferred(promote_stdmeasure(StdLogistic, StdLogistic)) === StdLogistic
+        @test @inferred(promote_stdmeasure(AnyStdMeasure, StdUniform)) === StdUniform
+        @test @inferred(promote_stdmeasure(StdUniform, AnyStdMeasure)) === StdUniform
+        @test @inferred(promote_stdmeasure(AnyStdMeasure, AnyStdMeasure)) === AnyStdMeasure
+        @test @inferred(promote_stdmeasure(NoStdTransport{Int}, StdNormal)) === NoStdTransport{Int}
+        @test @inferred(promote_stdmeasure(StdNormal, NoStdTransport{Int})) === NoStdTransport{Int}
+        @test @inferred(promote_stdmeasure(NoStdTransport{Int}, AnyStdMeasure)) === NoStdTransport{Int}
+        @test @inferred(promote_stdmeasure(AnyStdMeasure, NoStdTransport{Int})) === NoStdTransport{Int}
+        @test @inferred(promote_stdmeasure(StdUniform, StdExponential, AnyStdMeasure, StdLogistic)) === StdLogistic
+    end
 end

From 72f238e2703dd86a160f7c98a4eb6fed3b3dac56 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 17 Sep 2026 22:53:24 +0200
Subject: [PATCH 080/122] Evaluate densities of powers over flat variate
 storage

The batched density core unwraps nested powers once and evaluates the
point kernel of the base measure in one call over the flat variate
storage, with the variate dimensions leading, followed by the power and
batch dimensions. The power dimensions are summed afterwards. Batches and
variates given as nested arrays are fused at the entry point, nested
arrays without flat storage of a known layout are evaluated level by
level. Powers of measures with a known variate size route through
batched_logdensityof_impl, so batched kernels of the base measure apply
to powers as well.

The default batched kernel for scalar variates is a lazy broadcast, so
the sum for a single variate of a power stays allocation free. The
variadic power axes machinery, the static-zero-size and primitive-power
density specializations and the return type inference are gone.

Created by generative AI.
---
 src/MeasureBase.jl       |   3 +-
 src/combinators/power.jl |  28 +---
 src/density-batched.jl   | 335 +++++++++++++++++++++------------------
 test/logdensities.jl     |  55 ++++++-
 4 files changed, 237 insertions(+), 184 deletions(-)

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 52d6ec14..5bf5b4be 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -62,7 +62,8 @@ using HeterogeneousComputing: real_numtype
 
 using ArraysOfArrays:
     ArrayOfSimilarArrays, VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays,
-    VectorOfSimilarVectors, flatview
+    VectorOfSimilarVectors, flatview, fused, stacked, sliced, getsplitmode,
+    is_memordered_splitmode, AbstractSplitMode, UnknownSplitMode, NonSplitMode
 
 using OneTwoMany: firstarg, secondarg
 
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 443e4532..590e576d 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -100,22 +100,11 @@ params(d::PowerMeasure) = params(first(marginals(d)))
     basemeasure(d.parent)^d.axes
 end
 
-# Power structure is unwrapped into the power axes arguments of the batched
-# density machinery (see density-batched.jl), which fuses evaluation over
-# flat variate storage:
+# Densities of powers are evaluated by the batched density machinery over
+# the flat variate storage (see density-batched.jl):
 
-for head in [:logdensityof_impl, :logdensity_def]
-    @eval @inline function $head(d::PowerMeasure, x)
-        _powered_ld($head, pwr_base(d), x, pwr_axes(d))
-    end
-
-    @eval @inline function $head(
-        ::PowerMeasure{<:Any,<:Tuple{Vararg{StaticOneToLike{0}}}},
-        x,
-    )
-        static(0.0)
-    end
-end
+@inline logdensityof_impl(μ::PowerMeasure, x) = _powered_ld(logdensityof_impl, μ, x)
+@inline logdensity_def(μ::PowerMeasure, x) = _powered_ld(logdensity_def, μ, x)
 
 @inline function insupport(μ::PowerMeasure, x)
     p = μ.parent
@@ -165,13 +154,4 @@ end
 
 massof(m::PowerMeasure) = massof(m.parent)^prod(m.axes)
 
-logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0)
-
-# Disambiguation with the static-zero-size power density method:
-function logdensity_def(
-    ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0}}}},
-    ::Any,
-) where {P<:PrimitiveMeasure}
-    static(0.0)
-end
 
diff --git a/src/density-batched.jl b/src/density-batched.jl
index a29effd4..dd4910bf 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -8,192 +8,211 @@ export logdensities
 Compute the log-density of `μ` at each point in `X`.
 
 Returns an array of the shape of `X`, semantically equivalent to
-`logdensityof.(Ref(μ), X)`. The computation may be fused across points,
-though: power measures with flat variate storage (e.g. based on
-`ArraysOfArrays.ArrayOfSimilarArrays`) evaluate as a single flat broadcast
-plus a segmented reduction over the underlying flat data, compatible with
-GPU-backed storage.
-
-Measure types should specialize [`MeasureBase.batched_logdensityof_impl`](@ref)
-instead of `logdensities` itself.
+`logdensityof.(Ref(μ), X)`. Batches with flat storage of a known variate
+size (e.g. `ArraysOfArrays.ArrayOfSimilarArrays`) are evaluated in one
+fused operation over the flat data, compatible with GPU-backed storage.
+
+For measures with array-valued variates, `X` may also be the flat storage
+of the batch itself, with the variate dimensions leading (see
+[`MeasureBase.mspace_flatsize`](@ref)). The result then has the remaining
+dimensions of `X`.
+
+Measure types should specialize
+[`MeasureBase.batched_logdensityof_impl`](@ref) instead of `logdensities`
+itself.
 """
 function logdensities end
 
-@inline function logdensities(μ::AbstractMeasure, X::AbstractArray)
-    _logdensities(logdensityof_impl, μ, X)
-end
+@inline logdensities(μ::AbstractMeasure, X::AbstractArray) = _materialize(_batched_ld(logdensityof_impl, μ, X))
 
 """
-    MeasureBase.batched_logdensityof_impl(μ::AbstractMeasure, X::AbstractArray)
-
-Implements [`logdensities(μ, X)`](@ref logdensities) for arrays `X` of
-plain `μ`-variates. Power measures never reach `batched_logdensityof_impl`, their
-power structure is processed generically beforehand.
-
-Measure types that support fused multi-point evaluation should specialize
-`batched_logdensityof_impl`. Implementations must preserve the shape of `X` and
-must handle points outside the support of `μ` (the result must be `-Inf`
-at such points).
-
-The default implementation broadcasts the log-density over `X` for
-measures with scalar variates and falls back to a `map` over `X`
-otherwise.
+    MeasureBase.batched_logdensityof_impl(μ::AbstractMeasure, A::AbstractArray)
+
+Implements [`logdensities`](@ref) for a batch `A` of variates of `μ` in
+flat storage: the leading dimensions of `A` are the variate dimensions
+(see [`MeasureBase.mspace_flatsize`](@ref)), any further dimensions are
+batch dimensions. Returns the log-densities as an array over the batch
+dimensions, or a scalar if there are none.
+
+Power measures never reach `batched_logdensityof_impl`, their power
+structure is unwrapped beforehand. Implementations must handle points
+outside the support of `μ` (the result must be `-Inf` there).
+
+The default implementation broadcasts the log-density over `A` for
+measures with scalar variates and maps it over the variate slices of `A`
+otherwise. The result may be a lazy broadcast, callers materialize it
+where necessary.
 """
 function batched_logdensityof_impl end
 
-function batched_logdensityof_impl(μ::AbstractMeasure, X::AbstractArray)
-    _logdensities_generic(logdensityof_impl, μ, X)
+@inline function batched_logdensityof_impl(μ::AbstractMeasure, A::AbstractArray)
+    _batched_ld_generic(logdensityof_impl, μ, A)
 end
 
 # Batched density machinery, parameterized over the point-level density
 # function `f` (`logdensityof_impl` or `logdensity_def`).
 #
-# `_logdensities(f, μ, X, powers...)` treats each element of `X` as a
-# variate of `(μ^pN)^…^p1` for `powers = (p1, …, pN)` (power axes ordered
-# outermost first, i.e. in the order in which they are encountered when
-# descending into a variate) and returns the density sum for each element,
-# preserving the shape of `X`. Power measures are unwrapped into the power
-# axes arguments before any other dispatch happens, so implementation
-# methods only ever dispatch on plain measure types.
-
-@inline function _logdensities(f::F, μ, X::AbstractArray, powers::Vararg{Any,N}) where {F,N}
-    _logdensities_stripped(f, μ, X, powers...)
-end
-
-@inline function _logdensities(
-    f::F,
-    μ::PowerMeasure,
-    X::AbstractArray,
-    powers::Vararg{Any,N},
-) where {F,N}
-    _logdensities(f, pwr_base(μ), X, powers..., pwr_axes(μ))
-end
-
-@inline function _logdensities_stripped(f::F, μ, X::AbstractArray) where {F}
-    _batched_logdensityof_impl(f, μ, X)
-end
-
-function _logdensities_stripped(
-    f::F,
-    μ,
-    X::AbstractArray,
-    p1,
-    powers::Vararg{Any,N},
-) where {F,N}
-    map(x -> _powered_ld(f, μ, x, p1, powers...), X)
-end
-
-function _logdensities_stripped(
-    f::F,
-    μ,
-    X::ArrayOfSimilarArrays{<:Number},
-    p1,
-    powers::Vararg{Any,N},
-) where {F,N}
-    _logdensities_fused(f, μ, X, mspace_elsize(μ), p1, powers...)
-end
-
-# Absolute densities go through the `batched_logdensityof_impl` extension point:
-@inline function _batched_logdensityof_impl(::typeof(logdensityof_impl), μ, X::AbstractArray)
-    batched_logdensityof_impl(μ, X)
-end
-
-@inline function _batched_logdensityof_impl(f::F, μ, X::AbstractArray) where {F}
-    _logdensities_generic(f, μ, X)
-end
-
-@inline function _logdensities_generic(f::F, μ, X::AbstractArray) where {F}
-    _logdensities_byelsize(f, μ, X, mspace_elsize(μ))
-end
-
-# Scalar variates evaluate as a single flat broadcast:
-@inline function _logdensities_byelsize(
-    f::F,
-    μ,
-    X::AbstractArray{<:Number},
-    ::Tuple{},
-) where {F}
-    broadcast(Base.Fix1(f, μ), X)
-end
-
-@inline function _logdensities_byelsize(f::F, μ, X::AbstractArray, ::Any) where {F}
-    map(Base.Fix1(f, μ), X)
-end
-
-# Scalar-variate measure with flat variate storage: evaluate as a single
-# flat broadcast followed by a segmented reduction over the per-point
-# power structure:
-function _logdensities_fused(
-    f::F,
-    μ,
-    X::ArrayOfSimilarArrays{<:Number,M},
-    ::Tuple{},
-    powers::Vararg{Any,N},
-) where {F,M,N}
-    sz_inner = _flat_powers_size(powers...)
-    if length(sz_inner) == M
-        X_flat = flatview(X)
-        if ntuple(i -> size(X_flat, i), Val(M)) != sz_inner
-            throw(ArgumentError("Size of variates doesn't match size of power measure"))
-        end
-        ld_flat = broadcast(Base.Fix1(f, μ), X_flat)
-        reshape(sum(ld_flat, dims = ntuple(identity, Val(M))), size(X))
-    else
-        map(x -> _powered_ld(f, μ, x, powers...), X)
-    end
+# Variates and batches with flat storage are evaluated in one call of the
+# batched point kernel of the base measure: the leading dimensions of the
+# flat data are the variate dimensions of the base measure, followed by the
+# power dimensions and any batch dimensions. The power dimensions are summed
+# afterwards. Nested arrays without flat storage of a known layout are
+# evaluated level by level.
+
+struct NoFlatStorage end
+
+@inline _batched_ld(f::F, μ, X::AbstractArray) where {F} = _batched_ld_sized(f, μ, X, mspace_flatsize(μ))
+
+@inline function _batched_ld_sized(f::F, μ, X::AbstractArray, sz_flat::SizeLike) where {F}
+    _batched_ld_flat(f, μ, X, _flat_storage(X), sz_flat)
+end
+
+@inline function _batched_ld_sized(f::F, μ, X::AbstractArray, ::NoMSpaceElementSize) where {F}
+    map(x -> _pointwise_ld(f, μ, x), X)
+end
+
+@inline function _batched_ld_flat(f::F, μ, X, X_flat::AbstractArray, sz_flat) where {F}
+    ν, n_pwr = _pwr_unwrap(μ)
+    _check_flatsize(X_flat, sz_flat)
+    _sum_leading_dims(_batched_kernel(f, ν, X_flat), n_pwr)
 end
 
-function _logdensities_fused(
-    f::F,
-    μ,
-    X::AbstractArray,
-    ::Any,
-    p1,
-    powers::Vararg{Any,N},
-) where {F,N}
-    map(x -> _powered_ld(f, μ, x, p1, powers...), X)
+@inline function _batched_ld_flat(f::F, μ, X, ::NoFlatStorage, sz_flat) where {F}
+    map(x -> _pointwise_ld(f, μ, x), X)
 end
 
-# The flat size of a variate of `(μ^pN)^…^p1` for a scalar-variate `μ`,
-# innermost power axes vary fastest:
-@inline _flat_powers_size() = ()
-@inline function _flat_powers_size(p1, powers::Vararg{Any,N}) where {N}
-    (_flat_powers_size(powers...)..., axes2size(p1)...)
+@inline _pointwise_ld(f::F, μ, x) where {F} = f(μ, x)
+@inline _pointwise_ld(f::F, μ::PowerMeasure, x) where {F} = _powered_ld(f, μ, x)
+
+# Log-density of a power measure at a single variate:
+
+@inline _powered_ld(f::F, μ::PowerMeasure, x) where {F} = _powered_ld_sized(f, μ, x, mspace_flatsize(μ))
+
+@inline function _powered_ld_sized(f::F, μ::PowerMeasure, x, sz_flat::SizeLike) where {F}
+    _powered_ld_flat(f, μ, x, _flat_storage(x), sz_flat)
 end
 
-# Scalar counterpart of `_logdensities`: log-density of `(μ^pN)^…^p1` at a
-# single variate `x`.
-@inline _powered_ld(f::F, μ, x) where {F} = f(μ, x)
+@inline function _powered_ld_sized(f::F, μ::PowerMeasure, x, ::NoMSpaceElementSize) where {F}
+    _powered_ld_pointwise(f, μ, x)
+end
 
-@inline function _powered_ld(
-    f::F,
-    μ::PowerMeasure,
-    x,
-    p1,
-    powers::Vararg{Any,N},
-) where {F,N}
-    _powered_ld(f, pwr_base(μ), x, p1, powers..., pwr_axes(μ))
+@inline function _powered_ld_flat(f::F, μ::PowerMeasure, x, x_flat::AbstractArray, sz_flat) where {F}
+    ν, n_pwr = _pwr_unwrap(μ)
+    _check_flatsize(x_flat, sz_flat)
+    _sum_leading_dims(_batched_kernel(f, ν, x_flat), n_pwr)
 end
 
-@inline function _powered_ld(f::F, μ, x, p1, powers::Vararg{Any,N}) where {F,N}
-    if axes2size(p1) != maybestatic_size(x)
+@inline _powered_ld_flat(f::F, μ::PowerMeasure, x, ::NoFlatStorage, sz_flat) where {F} = _powered_ld_pointwise(f, μ, x)
+
+# Sum of the point-level densities over the elements of the variate:
+@inline function _powered_ld_pointwise(f::F, μ::PowerMeasure, x::AbstractArray) where {F}
+    if maybestatic_size(x) != pwr_size(μ)
         throw(ArgumentError("Size of variate doesn't match size of power measure"))
     end
-    R = _powered_ld_type(f, μ, x, powers...)
-    if isempty(x)
-        zero(R)::R
-    else
-        # Conversion needed since summation can turn static into dynamic values:
-        convert(R, _powered_ld_sum(f, μ, x, powers...))::R
+    sum(Base.Fix1(_pointwise_ld_dyn, (f, pwr_base(μ))), x)
+end
+
+function _powered_ld_pointwise(f::F, ::PowerMeasure, x) where {F}
+    throw(ArgumentError("Variate of a power measure must be an array"))
+end
+
+@inline _pointwise_ld_dyn((f, μ), x) = dynamic(_pointwise_ld(f, μ, x))
+
+@inline _pwr_unwrap(μ) = (μ, static(0))
+@inline function _pwr_unwrap(μ::PowerMeasure)
+    ν, n = _pwr_unwrap(pwr_base(μ))
+    ν, n + static(length(pwr_axes(μ)))
+end
+
+# Flat storage of a (nested) variate or batch: the underlying array of
+# memory-ordered split arrays, a stacked copy for other known split modes.
+# Nested arrays of unknown layout have no flat storage.
+
+@inline _flat_storage(x::AbstractArray{<:Number}) = x
+@inline _flat_storage(x::AbstractArray) = _flat_storage_bymode(x, getsplitmode(x))
+@inline _flat_storage(x) = NoFlatStorage()
+
+@inline function _flat_storage_bymode(x::AbstractArray, smode::AbstractSplitMode)
+    _flat_storage(is_memordered_splitmode(smode) ? fused(x) : stacked(x))
+end
+@inline _flat_storage_bymode(::AbstractArray, ::UnknownSplitMode) = NoFlatStorage()
+@inline _flat_storage_bymode(::AbstractArray, ::NonSplitMode) = NoFlatStorage()
+
+@inline function _check_flatsize(A::AbstractArray, sz_flat::SizeLike)
+    n = length(sz_flat)
+    if ndims(A) < n || ntuple(i -> size(A, i), Val(n)) != Tuple(sz_flat)
+        throw(ArgumentError("Size of variate doesn't match size of measure"))
     end
+    return nothing
+end
+
+# Point kernel of the base measure over the flat batch:
+
+@inline _batched_kernel(::typeof(logdensityof_impl), ν, A::AbstractArray) = batched_logdensityof_impl(ν, A)
+@inline _batched_kernel(f::F, ν, A::AbstractArray) where {F} = _batched_ld_generic(f, ν, A)
+
+# Powers of primitive measures have log-density zero relative to their base:
+@inline function _batched_kernel(::typeof(logdensity_def), ν::PrimitiveMeasure, A::AbstractArray)
+    FillArrays.Zeros{Float64}(size(A))
+end
+
+@inline _batched_ld_generic(f::F, ν, A::AbstractArray) where {F} = _batched_ld_byflatsize(f, ν, A, mspace_flatsize(ν))
+
+# Scalar variates: one lazy broadcast over the whole batch, so that
+# reductions over it don't need to allocate the intermediate result.
+# Static results of point kernels are made dynamic, to keep reductions
+# over them type stable.
+@inline function _batched_ld_byflatsize(f::F, ν, A::AbstractArray, ::Tuple{}) where {F}
+    Broadcast.instantiate(Broadcast.broadcasted(dynamic ∘ Base.Fix1(f, ν), A))
+end
+
+# Array variates: map over the variate slices, or evaluate directly if `A`
+# is a single variate.
+@inline function _batched_ld_byflatsize(f::F, ν, A::AbstractArray, sz::SizeLike) where {F}
+    _batched_ld_slices(f, ν, A, Val(length(sz)))
 end
 
-@inline _powered_ld_sum(f::F, μ, x) where {F} = sum(Base.Fix1(f, μ), x)
+# Variates of unknown size: the elements of `A` are the variates.
+@inline function _batched_ld_byflatsize(f::F, ν, A::AbstractArray, ::NoMSpaceElementSize) where {F}
+    map(Base.Fix1(f, ν), A)
+end
 
-@inline function _powered_ld_sum(f::F, μ, x, p1, powers::Vararg{Any,N}) where {F,N}
-    sum(_logdensities(f, μ, x, p1, powers...))
+@inline _batched_ld_slices(f::F, ν, A::AbstractArray{<:Any,N}, ::Val{N}) where {F,N} = f(ν, A)
+@inline function _batched_ld_slices(f::F, ν, A::AbstractArray, ::Val{M}) where {F,M}
+    map(Base.Fix1(f, ν), sliced(A, Val(M)))
 end
 
-@inline function _powered_ld_type(f::F, ::MU, x, powers::Vararg{Any,N}) where {F,MU,N}
-    Core.Compiler.return_type(_powered_ld, Tuple{F,MU,eltype(x),map(typeof, powers)...})
+# Sum over the leading `N` dimensions; a full reduction yields a scalar.
+# Lazy broadcasts are reduced without materialization where the broadcast
+# style supports it, and materialized before partial reductions.
+
+const _LazyBroadcast = Broadcast.Broadcasted
+const _EagerReducibleBroadcast = Broadcast.Broadcasted{<:Union{Broadcast.DefaultArrayStyle,StaticArrays.StaticArrayStyle}}
+
+@inline _materialize(bc::_LazyBroadcast) = copy(bc)
+@inline _materialize(x) = x
+
+@inline _sum_leading_dims(x::Number, ::StaticInteger{0}) = x
+@inline _sum_leading_dims(A::AbstractArray, n::StaticInteger) = _sum_leading_dims_impl(A, n, static(ndims(A)))
+@inline _sum_leading_dims(bc::_LazyBroadcast, n::StaticInteger) = _sum_leading_dims_lazy(bc, n, static(ndims(bc)))
+@inline _sum_leading_dims_lazy(bc::_LazyBroadcast, ::StaticInteger{0}, ::StaticInteger) = bc
+@inline _sum_leading_dims_lazy(bc::_LazyBroadcast, ::StaticInteger{0}, ::StaticInteger{0}) = bc
+@inline _sum_leading_dims_lazy(bc::_EagerReducibleBroadcast, ::StaticInteger{0}, ::StaticInteger{0}) = bc
+@inline function _sum_leading_dims_lazy(bc::_EagerReducibleBroadcast, ::StaticInteger{N}, ::StaticInteger{N}) where {N}
+    # Empty broadcasts have no known element type to reduce over lazily:
+    isempty(bc) ? sum(copy(bc)) : sum(bc)
+end
+@inline _sum_leading_dims_lazy(bc::_LazyBroadcast, ::StaticInteger{N}, ::StaticInteger{N}) where {N} = sum(copy(bc))
+@inline function _sum_leading_dims_lazy(bc::_LazyBroadcast, n::StaticInteger, ::StaticInteger)
+    _sum_leading_dims(copy(bc), n)
+end
+@inline _sum_leading_dims_impl(A::AbstractArray, ::StaticInteger{0}, ::StaticInteger) = A
+@inline _sum_leading_dims_impl(A::AbstractArray, ::StaticInteger{0}, ::StaticInteger{0}) = A
+@inline _sum_leading_dims_impl(A::AbstractArray, ::StaticInteger{N}, ::StaticInteger{N}) where {N} = sum(A)
+@inline function _sum_leading_dims_impl(A::AbstractArray, ::StaticInteger{N}, ::StaticInteger) where {N}
+    dropdims(_sum_dims_seq(A, static(N)); dims = ntuple(identity, Val(N)))
+end
+@inline _sum_dims_seq(A::AbstractArray, ::StaticInteger{0}) = A
+@inline function _sum_dims_seq(A::AbstractArray, ::StaticInteger{N}) where {N}
+    _sum_dims_seq(sum(A; dims = N), static(N - 1))
 end
diff --git a/test/logdensities.jl b/test/logdensities.jl
index 2f68f72f..e823a138 100644
--- a/test/logdensities.jl
+++ b/test/logdensities.jl
@@ -3,8 +3,10 @@
 using Test
 
 using MeasureBase
-using MeasureBase: logdensities, StdNormal, StdUniform
+using MeasureBase: logdensities, logdensity_def, StdNormal, StdUniform, Dirac, Lebesgue, LebesgueBase, superpose, weightedmeasure
 using ArraysOfArrays: VectorOfSimilarVectors, sliced, flatview
+using StaticArrays: SVector, @SVector, @SMatrix
+using Static: static
 using IrrationalConstants: log2π
 import JLArrays
 using JLArrays: JLArray
@@ -52,6 +54,14 @@ stdnormal_ld(x) = -(x^2 + log2π) / 2
         @test logdensities(mprod, X) ≈ logdensityof.(Ref(mprod), X)
     end
 
+    @testset "unknown variate size" begin
+        mix = superpose(weightedmeasure(log(0.3), StdNormal()), weightedmeasure(log(0.7), StdUniform()))
+        X = randn(4, 5)
+        @test logdensities(mix, X) ≈ logdensityof.(Ref(mix), X)
+        @test logdensityof(mix^4, X[:, 1]) ≈ sum(logdensityof.(Ref(mix), X[:, 1]))
+        @test logdensities(mix^4, sliced(X, 1)) ≈ vec(sum(logdensityof.(Ref(mix), X), dims = 1))
+    end
+
     @testset "size mismatch" begin
         @test_throws ArgumentError logdensities(StdNormal()^3, [randn(3), randn(2)])
         @test_throws ArgumentError logdensities(
@@ -60,6 +70,49 @@ stdnormal_ld(x) = -(x^2 + log2π) / 2
         )
     end
 
+    @testset "flat batch storage and array-variate bases" begin
+        m3 = StdNormal()^3
+        Xf = randn(3, 10)
+        @test @inferred(logdensities(m3, Xf)) ≈ vec(sum(stdnormal_ld.(Xf), dims = 1))
+        x = randn(3)
+        @test @inferred(logdensities(m3, x)) ≈ sum(stdnormal_ld, x)
+
+        mpp = (StdNormal()^(2, 3))^4
+        Xpp_flat = randn(2, 3, 4, 7)
+        @test @inferred(logdensities(mpp, Xpp_flat)) ≈ vec(sum(stdnormal_ld.(Xpp_flat), dims = (1, 2, 3)))
+        Xpp_nested = sliced(sliced(Xpp_flat, 2), 1)
+        @test logdensities(mpp, Xpp_nested) ≈ logdensities(mpp, Xpp_flat)
+        xpp = randn(2, 3, 4)
+        @test @inferred(logdensityof(mpp, xpp)) ≈ logdensityof(mpp, [xpp[:, :, i] for i in 1:4])
+        @test logdensityof(mpp, sliced(xpp, 2)) ≈ logdensityof(mpp, xpp)
+
+        mvec = Dirac([1.0, 2.0])^3
+        @test @inferred(logdensities(mvec, [fill([1.0, 2.0], 3) for _ in 1:2])) == [0.0, 0.0]
+    end
+
+    @testset "static variates" begin
+        m3 = StdNormal()^static(3)
+        xs = @SVector randn(3)
+        f(x) = logdensityof(m3, x)
+        @test @inferred(f(xs)) ≈ sum(stdnormal_ld, xs)
+        @test @allocated(f(xs)) == 0
+        g(x) = logdensityof(StdNormal()^3, x)
+        xd = randn(3)
+        @test @inferred(g(xd)) ≈ sum(stdnormal_ld, xd)
+        @test @allocated(g(xd)) == 0
+        Xs = @SMatrix randn(3, 4)
+        @test @inferred(logdensities(m3, Xs)) ≈ vec(sum(stdnormal_ld.(Xs), dims = 1))
+        @test logdensities(m3, Xs) isa SVector{4}
+        @test @inferred(logdensityof(StdNormal()^static(0), SVector{0,Float64}())) == 0
+        @test @inferred(logdensityof(StdNormal()^0, Float64[])) == 0
+    end
+
+    @testset "powers of primitive measures" begin
+        @test @inferred(logdensity_def(Lebesgue()^3, randn(3))) == 0
+        @test @inferred(logdensity_def(LebesgueBase()^(2, 2), randn(2, 2))) == 0
+        @test @inferred(logdensityof(Lebesgue()^3, randn(3))) == 0
+    end
+
     @testset "GPU array semantics" begin
         JLArrays.allowscalar(false)
 

From c454330ca843ce91462694dcd5b1746565e88618 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 17 Sep 2026 23:37:57 +0200
Subject: [PATCH 081/122] Make density evaluation branch-free and number-type
 preserving

Support checks enter density evaluation as masks: relative densities
evaluate the kernels unconditionally and select the result with ifelse,
superpositions compute their density in log space with masked components
and a branch-free logsumexp, density measures, spike mixtures, products
and restricted measures no longer branch on support values, and
NoFastInsupport means unconditional evaluation.

Log-density kernels return numbers of the number type of the variate:
primitive measures return typed zeros instead of static ones, which also
keeps Mooncake from hashing static constants, and plain floating-point
log-weights adopt the variate number type while weights that carry
derivatives promote as before.

Created by generative AI.
---
 ext/MeasureBaseChainRulesCoreExt.jl |  9 ---
 ext/MeasureBaseMooncakeExt.jl       |  3 -
 src/combinators/product.jl          | 18 ++----
 src/combinators/restricted.jl       |  2 +-
 src/combinators/spikemixture.jl     | 22 ++------
 src/combinators/superpose.jl        | 49 ++++++++--------
 src/combinators/weighted.jl         | 12 ++--
 src/density-core.jl                 | 88 ++++++++++-------------------
 src/density.jl                      | 22 +++-----
 src/primitive.jl                    | 11 ++--
 src/primitives/dirac.jl             |  4 +-
 src/primitives/lebesgue.jl          |  6 +-
 src/utils.jl                        |  4 +-
 test/numtype.jl                     | 39 +++++++++++++
 test/runtests.jl                    |  1 +
 test/test_mooncake.jl               |  2 -
 16 files changed, 134 insertions(+), 158 deletions(-)
 create mode 100644 test/numtype.jl

diff --git a/ext/MeasureBaseChainRulesCoreExt.jl b/ext/MeasureBaseChainRulesCoreExt.jl
index 25019da1..1f457026 100644
--- a/ext/MeasureBaseChainRulesCoreExt.jl
+++ b/ext/MeasureBaseChainRulesCoreExt.jl
@@ -97,13 +97,4 @@ ChainRulesCore.rrule(::typeof(check_dof), ν, μ) = check_dof(ν, μ), _check_do
 _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
diff --git a/ext/MeasureBaseMooncakeExt.jl b/ext/MeasureBaseMooncakeExt.jl
index e401cb7d..a834bbec 100644
--- a/ext/MeasureBaseMooncakeExt.jl
+++ b/ext/MeasureBaseMooncakeExt.jl
@@ -8,7 +8,6 @@ using Mooncake: @zero_derivative, MinimalCtx
 
 using MeasureBase: isneginf, isposinf, _adignore_call
 using MeasureBase: check_dof, require_insupport, _origin_depth
-using MeasureBase: logdensityof_rt
 
 # Unlike Zygote, Mooncake differentiates the collection utilities
 # (`_pushfront`, etc., mutating code in general), `checked_arg` and
@@ -24,6 +23,4 @@ using MeasureBase: logdensityof_rt
 @zero_derivative MinimalCtx Tuple{typeof(_origin_depth),Any}
 @zero_derivative MinimalCtx Tuple{typeof(check_dof),Any,Any}
 
-@zero_derivative MinimalCtx Tuple{typeof(logdensityof_rt),Any,Any}
-
 end # module MeasureBaseMooncakeExt
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 22dae865..9d44112b 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -239,24 +239,16 @@ function _rand(rng::AbstractRNG, ::Type{T}, d::ProductMeasure, mar::AbstractArra
 end
 
 @inline function insupport(d::AbstractProductMeasure, x::AbstractArray)
-    mar = marginals(d)
-    # We might get lucky and know statically that everything is inbounds
-    T = Core.Compiler.return_type(insupport, Tuple{eltype(mar),eltype(x)})
-    T <: True || all(zip(x, mar)) do (xj, mj)
-        insupport(mj, xj) == true
-    end
+    _all_insupport(broadcast(_insupport_bool ∘ insupport, marginals(d), x))
 end
 
 @inline function insupport(d::AbstractProductMeasure, x)
-    for (mj, xj) in zip(marginals(d), x)
-        insup = dynamic(insupport(mj, xj))
-        if insup isa NoFastInsupport || insup == false
-            return insup
-        end
-    end
-    return true
+    mapreduce(insupport, _insupport_and, marginals(d), x)
 end
 
+@inline _all_insupport(A::AbstractArray{<:NoFastInsupport{T}}) where {T} = NoFastInsupport{T}()
+@inline _all_insupport(A::AbstractArray) = all(A)
+
 getdof(d::AbstractProductMeasure) = sum(getdof, marginals(d))
 fast_dof(d::AbstractProductMeasure) = sum(fast_dof, marginals(d))
 
diff --git a/src/combinators/restricted.jl b/src/combinators/restricted.jl
index 063c7f33..792d3a5b 100644
--- a/src/combinators/restricted.jl
+++ b/src/combinators/restricted.jl
@@ -10,7 +10,7 @@ end
 
 basemeasure(μ::RestrictedMeasure) = μ.base
 
-insupport(μ::RestrictedMeasure, x) = μ.predicate(x) && insupport(μ.base, x)
+insupport(μ::RestrictedMeasure, x) = _insupport_and(μ.predicate(x), insupport(μ.base, x))
 
 function Pretty.quoteof(d::RestrictedMeasure)
     qf = Pretty.quoteof(d.predicate)
diff --git a/src/combinators/spikemixture.jl b/src/combinators/spikemixture.jl
index e39d4230..216d66f9 100644
--- a/src/combinators/spikemixture.jl
+++ b/src/combinators/spikemixture.jl
@@ -23,23 +23,9 @@ end
 
 for func in [:logdensityof, :logdensity_def]
     @eval @inline function $func(μ::SpikeMixture, x)
-        # NOTE: We could instead write this as 
-        # R1 = typeof(log(one(μ.s))) 
-        # R2 = typeof(log(one(μ.w))) 
-
-        # which would rely on constant propagation insteadof type inference.
-        # We'll try this for now and come back to the question if we see
-        # problems.
-
-        R1 = Core.Compiler.return_type(log, Tuple{typeof(μ.s)})
-        R2 = Core.Compiler.return_type(log, Tuple{typeof(μ.w)})
-        R3 = Core.Compiler.return_type($func, Tuple{typeof(μ.m),typeof(x)})
-        R = promote_type(R1, R2, R3)
-        if iszero(x)
-            return convert(R, log(μ.s))::R
-        else
-            return convert(R, log(μ.w) + $func(μ.m, x))::R
-        end
+        ℓ_spike = dynamic(log(μ.s))
+        ℓ_parent = dynamic(log(μ.w)) + dynamic($func(μ.m, x))
+        ifelse(iszero(x), oftype(ℓ_parent, ℓ_spike), ℓ_parent)
     end
 end
 
@@ -53,4 +39,4 @@ end
 
 testvalue(::Type{T}, μ::SpikeMixture) where {T} = zero(T)
 
-insupport(μ::SpikeMixture, x) = dynamic(insupport(μ.m, x)) || iszero(x)
+insupport(μ::SpikeMixture, x) = _insupport_mask(insupport(μ.m, x)) | iszero(x)
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index d9b53bd5..369cbd5d 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -65,37 +65,44 @@ function Base.:+(μ::AbstractMeasure, ν::AbstractMeasure)
     superpose(μ, ν)
 end
 
-@inline _ulogexp(x) = exp(ULogarithmic, dynamic(x))
+# Masks components outside of their support with -Inf:
+@inline _masked_logd(ℓ, ins) = ifelse(_insupport_mask(ins), ℓ, oftype(ℓ, -Inf))
+
+# Branch-free logsumexp over the components, valid for infinite entries:
+@inline function _logsumexp_components(ℓs)
+    m = reduce(max, ℓs)
+    m_finite = ifelse(isfinite(m), m, zero(m))
+    m_finite + log(sum(map(ℓ -> exp(ℓ - m_finite), ℓs)))
+end
 
-function density_def(s::SuperpositionMeasure, x)
+# The density of a superposition relative to the superposition of the
+# component base measures, in log space: each component contributes its
+# own density, divided by the density of the superposed base measures
+# relative to its own base measure.
+function logdensity_def(s::SuperpositionMeasure, x)
     cs = values(s.components)
     αs = map(basemeasure, cs)
-    idxs = eachindex(cs)
-    sum(idxs) do i
-        dμᵢ_dαᵢ = _ulogexp(logdensity_def(cs[i], x))
-        istrue(insupport(cs[i], x)) || return zero(dμᵢ_dαᵢ)
-        dΣα_dαᵢ = sum(idxs) do j
-            dαⱼ_dαᵢ = _ulogexp(logdensity_rel(αs[j], αs[i], x))
-            istrue(insupport(cs[j], x)) ? dαⱼ_dαᵢ : zero(dαⱼ_dαᵢ)
-        end
-        dμᵢ_dαᵢ / dΣα_dαᵢ
+    terms = map(cs, αs) do cᵢ, αᵢ
+        ℓᵢ = _dynamic_logd(logdensity_def(cᵢ, x), x)
+        log_dΣα_dαᵢ = _logsumexp_components(map(cs, αs) do cⱼ, αⱼ
+            _masked_logd(logdensity_rel(αⱼ, αᵢ, x), insupport(cⱼ, x))
+        end)
+        _masked_logd(ℓᵢ - log_dΣα_dαᵢ, insupport(cᵢ, x))
     end
+    _logsumexp_components(terms)
 end
 
 @inline function logdensity_rel_def(μ::T, ν::T, x) where {T<:SuperpositionMeasure}
-    if μ === ν
-        return zero(return_type(logdensity_def, (μ, x)))
-    else
-        return logdensity_def(μ, x) - logdensity_def(ν, x)
-    end
+    ℓ = logdensity_def(μ, x) - logdensity_def(ν, x)
+    ifelse(μ === ν, zero(ℓ), ℓ)
 end
 
 function _superpos_logdensity_rel(s::SuperpositionMeasure, β, x)
     cs = values(s.components)
     ds = map(cs) do μ
-        istrue(insupport(μ, x)) ? dynamic(logdensity_rel(μ, β, x)) : -Inf
+        _masked_logd(logdensity_rel(μ, β, x), insupport(μ, x))
     end
-    logsumexp(ds)
+    _logsumexp_components(ds)
 end
 
 @inline logdensity_rel_def(s::SuperpositionMeasure, β, x) = _superpos_logdensity_rel(s, β, x)
@@ -105,7 +112,7 @@ end
 
 @inline logdensity_rel_def(s, β::SuperpositionMeasure, x) = -_superpos_logdensity_rel(β, s, x)
 
-@inline logdensity_def(s::SuperpositionMeasure, x) = log(density_def(s, x))
+@inline density_def(s::SuperpositionMeasure, x) = exp(logdensity_def(s, x))
 
 function basemeasure(μ::SuperpositionMeasure{<:Tuple})
     superpose(map(basemeasure, μ.components)...)
@@ -135,7 +142,5 @@ function Base.rand(rng::AbstractRNG, ::Type{T}, μ::SuperpositionMeasure) where
 end
 
 @inline function insupport(d::SuperpositionMeasure, x)
-    any(d.components) do c
-        dynamic(insupport(c, x))
-    end
+    mapreduce(c -> _insupport_mask(insupport(c, x)), |, values(d.components))
 end
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index dfa28e74..aa3feb3d 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -12,14 +12,18 @@ abstract type AbstractWeightedMeasure <: AbstractMeasure end
 # By default the weight for all measure is 1
 _logweight(::AbstractMeasure) = 0
 
-@inline function logdensity_def(d::AbstractWeightedMeasure, _)
-    d.logweight
-end
+@inline logdensity_def(d::AbstractWeightedMeasure, x) = _logweight_for(d.logweight, x)
+
+# Plain floating-point log-weights adopt the number type of the variate,
+# log-weights that carry more information (dual numbers, traced values)
+# promote as usual:
+@inline _logweight_for(w, x) = w
+@inline _logweight_for(w::Union{AbstractFloat,StaticFloat64}, x) = _logd_numtype(x)(dynamic(w))
 
 # The weight-shifted density of a support-safe base density is support-safe,
 # no explicit support check required:
 @inline function logdensityof_impl(d::AbstractWeightedMeasure, x)
-    d.logweight + logdensityof_impl(basemeasure(d), x)
+    _logweight_for(d.logweight, x) + logdensityof_impl(basemeasure(d), x)
 end
 
 function Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractWeightedMeasure) where {T}
diff --git a/src/density-core.jl b/src/density-core.jl
index 9a6b0649..181f147c 100644
--- a/src/density-core.jl
+++ b/src/density-core.jl
@@ -45,16 +45,33 @@ a [`MeasureBase.NoFastInsupport`](@ref)) and computes the density via
 [`unsafe_logdensityof`](@ref).
 """
 @inline function logdensityof_impl(μ::AbstractMeasure, x)
-    result = dynamic(unsafe_logdensityof(μ, x))
+    result = _dynamic_logd(unsafe_logdensityof(μ, x), x)
     _checksupport(insupport(μ, x), result)
 end
 
-@inline function logdensityof_rt(::T, ::U) where {T,U}
-    Core.Compiler.return_type(logdensityof, Tuple{T,U})
-end
+# Log-density kernels return numbers of the number type of the variate,
+# never static numbers, so that automatic differentiation and tracing see
+# ordinary floating point values throughout:
+@inline _logd_numtype(x) = float(real_numtype(typeof(x)))
+@inline _dynamic_logd(ℓ, x) = dynamic(ℓ)
+@inline _neg_inf_logd(x) = _logd_numtype(x)(-Inf)
+
+# Support checks as masks: `NoFastInsupport` means the density is evaluated
+# unconditionally.
+@inline _insupport_mask(ins) = ins == true
+@inline _insupport_mask(::NoFastInsupport) = true
+
+# Support checks as booleans, keeping `NoFastInsupport`:
+@inline _insupport_bool(ins) = ins == true
+@inline _insupport_bool(ins::NoFastInsupport) = ins
+
+# Combining support checks of components, `NoFastInsupport` is absorbing:
+@inline _insupport_and(a, b) = _insupport_bool(a) & _insupport_bool(b)
+@inline _insupport_and(a::NoFastInsupport, ::Any) = a
+@inline _insupport_and(::Any, b::NoFastInsupport) = b
+@inline _insupport_and(a::NoFastInsupport, ::NoFastInsupport) = a
 
-_checksupport(cond, result) = ifelse(cond == true, result, oftype(result, -Inf))
-@inline _checksupport(::NoFastInsupport, result) = result
+@inline _checksupport(cond, result) = ifelse(_insupport_mask(cond), result, oftype(result, -Inf))
 
 """
     MeasureBase.logdensityof_with_rest(μ::AbstractMeasure, x)
@@ -153,9 +170,7 @@ end
         # if b_{i} isa typeof(b_{i - 1})
         #     return ℓ_{i - 1}
         # end
-        ℓ_{i} = let Δℓ_{i} = logdensity_def(b_{i}, x)
-            ℓ_{i - 1} + Δℓ_{i}
-        end
+        ℓ_{i} = ℓ_{i - 1} + logdensity_def(b_{i}, x)
     end
     return ℓ_10
 end
@@ -168,55 +183,12 @@ whether `x` is in the support of `m1` or `m2` (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`. 
 """
-@inline function logdensity_rel(μ::M, ν::N, x::X) where {M,N,X}
-    inμ = insupport(μ, x)
-    inν = insupport(ν, x)
-    return _logdensity_rel_impl(μ, ν, x, inμ, inν)
-end
-
-@inline function _logdensity_rel_impl(μ::M, ν::N, x::X, inμ::Bool, inν::Bool) where {M,N,X}
-    T = unstatic(
-        promote_type(
-            return_type(logdensity_def, (μ, x)),
-            return_type(logdensity_def, (ν, x)),
-        ),
-    )
-    istrue(inμ) || return convert(T, ifelse(inν, -Inf, NaN))
-    istrue(inν) || return convert(T, Inf)
-
-    return unsafe_logdensity_rel(μ, ν, x)
-end
-
-@inline function _logdensity_rel_impl(
-    μ::M,
-    ν::N,
-    x::X,
-    @nospecialize(::NoFastInsupport),
-    @nospecialize(::NoFastInsupport)
-) where {M,N,X}
-    unsafe_logdensity_rel(μ, ν, x)
-end
-
-@inline function _logdensity_rel_impl(
-    μ::M,
-    ν::N,
-    x::X,
-    inμ::Bool,
-    @nospecialize(::NoFastInsupport)
-) where {M,N,X}
-    logd = unsafe_logdensity_rel(μ, ν, x)
-    return istrue(inμ) ? logd : oftype(logd, -Inf)
-end
-
-@inline function _logdensity_rel_impl(
-    μ::M,
-    ν::N,
-    x::X,
-    @nospecialize(::NoFastInsupport),
-    inν::Bool
-) where {M,N,X}
-    logd = unsafe_logdensity_rel(μ, ν, x)
-    return istrue(inν) ? logd : oftype(logd, +Inf)
+@inline function logdensity_rel(μ, ν, x)
+    inμ = _insupport_mask(insupport(μ, x))
+    inν = _insupport_mask(insupport(ν, x))
+    logd = _dynamic_logd(unsafe_logdensity_rel(μ, ν, x), x)
+    outside = ifelse(inμ, oftype(logd, +Inf), ifelse(inν, oftype(logd, -Inf), oftype(logd, NaN)))
+    return ifelse(inμ & inν, logd, outside)
 end
 
 """
diff --git a/src/density.jl b/src/density.jl
index 67d34464..900332ae 100644
--- a/src/density.jl
+++ b/src/density.jl
@@ -95,7 +95,7 @@ struct DensityMeasure{F,B} <: AbstractMeasure
 end
 
 @inline function insupport(d::DensityMeasure, x)
-    insupport(d.base, x) == true && isfinite(logdensityof(getfield(d, :f), x))
+    _insupport_mask(insupport(d.base, x)) & isfinite(logdensityof(getfield(d, :f), x))
 end
 
 function Pretty.tile(μ::DensityMeasure{F,B}) where {F,B}
@@ -226,18 +226,10 @@ density_def(μ::DensityMeasure, x) = densityof(μ.f, x)
 
 function logdensityof_impl(μ::DensityMeasure, x::Any)
     integrand, μ_base = μ.f, μ.base
-
-    base_logval = logdensityof(μ_base, x)
-
-    T = typeof(base_logval)
-    U = logdensityof_rt(integrand, x)
-    R = promote_type(T, U)
-
-    # Don't evaluate base measure if integrand is zero or NaN
-    if isneginf(base_logval)
-        R(-Inf)
-    else
-        integrand_logval = logdensityof(integrand, x)
-        convert(R, integrand_logval + base_logval)::R
-    end
+    base_logval = dynamic(logdensityof(μ_base, x))
+    integrand_logval = dynamic(logdensityof(integrand, x))
+    logval = integrand_logval + base_logval
+    # Outside of the support of the base measure the integrand may be
+    # anything, including NaN:
+    ifelse(isneginf(base_logval), oftype(logval, -Inf), logval)
 end
diff --git a/src/primitive.jl b/src/primitive.jl
index 0a2aba37..9589c8ca 100644
--- a/src/primitive.jl
+++ b/src/primitive.jl
@@ -9,9 +9,9 @@ measures satisfy the following laws:
 
     basemeasure(μ::PrimitiveMeasure) = μ
 
-    logdensity_def(μ::PrimitiveMeasure, x) = 0.0
+    logdensity_def(μ::PrimitiveMeasure, x) == 0
 
-    logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0
+    logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} == 0
 """
 abstract type PrimitiveMeasure <: AbstractMeasure end
 
@@ -19,12 +19,11 @@ basemeasure(μ::PrimitiveMeasure) = μ
 
 @inline basemeasure_depth(::PrimitiveMeasure) = static(0)
 
-@inline logdensityof_impl(::PrimitiveMeasure, x::Number) = zero(float(typeof(x)))
-@inline logdensityof_impl(::PrimitiveMeasure, x) = static(0.0)
+@inline logdensityof_impl(::PrimitiveMeasure, x) = zero(_logd_numtype(x))
 
-logdensity_def(::PrimitiveMeasure, x) = static(0.0)
+logdensity_def(::PrimitiveMeasure, x) = zero(_logd_numtype(x))
 
-logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0
+logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = zero(_logd_numtype(x))
 
 function Pretty.quoteof(μ::M) where {M<:PrimitiveMeasure}
     :($M())
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 32d5d1de..3b134463 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -28,10 +28,10 @@ function logdensityof_impl(μ::Dirac, x::Number)
     _checksupport(insupport(μ, x), zero(R))
 end
 
-logdensityof_impl(μ::Dirac, x) = _checksupport(insupport(μ, x), 0.0)
+logdensityof_impl(μ::Dirac, x) = _checksupport(insupport(μ, x), zero(_logd_numtype(x)))
 
 logdensity_def(::Dirac, x::Number) = zero(float(typeof(x)))
-logdensity_def(::Dirac, x) = 0.0
+logdensity_def(::Dirac, x) = zero(_logd_numtype(x))
 
 Base.rand(::Random.AbstractRNG, T::Type, μ::Dirac) = μ.x
 
diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl
index 5846c765..7bd718cf 100644
--- a/src/primitives/lebesgue.jl
+++ b/src/primitives/lebesgue.jl
@@ -15,9 +15,9 @@ insupport(::LebesgueBase, x) = true
 
 insupport(::LebesgueBase) = Returns(true)
 
-logdensity_rel_def(::LebesgueBase, ::CountingBase, x) = -Inf
+logdensity_rel_def(::LebesgueBase, ::CountingBase, x) = _neg_inf_logd(x)
 
-logdensity_rel_def(::CountingBase, ::LebesgueBase, x) = Inf
+logdensity_rel_def(::CountingBase, ::LebesgueBase, x) = -_neg_inf_logd(x)
 
 @inline getdof(::LebesgueBase) = static(1)
 
@@ -52,7 +52,7 @@ gentype(::Lebesgue) = Float64
 Lebesgue() = Lebesgue(ℝ)
 
 @inline mspace_elsize(μ::Lebesgue) = _valueset_elsize(μ.support)
-@inline mspace_flatsize(μ::Lebesgue) = _valueset_elsize(μ.support)
+@inline mspace_flatsize(μ::Lebesgue) = _valueset_flatsize(μ.support)
 
 testvalue(::Type{T}, d::Lebesgue) where {T} = testvalue(T, d.support)::T
 
diff --git a/src/utils.jl b/src/utils.jl
index b6035e96..f73de9a4 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -140,8 +140,8 @@ fcomp(::typeof(identity), ::typeof(identity)) = identity
 
 near_neg_inf(::Type{T}) where {T<:Number} = T(-1E38) # Still fits into Float32
 
-isneginf(x) = isinf(x) && x < zero(x)
-isposinf(x) = isinf(x) && x > zero(x)
+isneginf(x) = isinf(x) & (x < zero(x))
+isposinf(x) = isinf(x) & (x > zero(x))
 
 isapproxzero(x::T) where {T<:Real} = x ≈ zero(T)
 isapproxzero(A::AbstractArray) = all(isapproxzero, A)
diff --git a/test/numtype.jl b/test/numtype.jl
new file mode 100644
index 00000000..5a962ae9
--- /dev/null
+++ b/test/numtype.jl
@@ -0,0 +1,39 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using MeasureBase
+using MeasureBase: StdNormal, StdUniform, LebesgueBase, Lebesgue, Dirac
+using MeasureBase: logdensity_def, logdensity_rel, weightedmeasure, superpose, logdensities, mbind
+using Static: static
+import ForwardDiff
+
+@testset "number type of log-densities" begin
+    x = 0.5f0
+    xf = randn(Float32, 3)
+    uf = rand(Float32, 3)
+    mix = superpose(weightedmeasure(log(0.3f0), StdNormal()), weightedmeasure(log(0.7f0), StdUniform()))
+
+    @test @inferred(logdensityof(StdNormal(), x)) isa Float32
+    @test @inferred(logdensityof(StdNormal()^3, xf)) isa Float32
+    @test @inferred(MeasureBase.unsafe_logdensityof(StdNormal()^3, xf)) isa Float32
+    @test @inferred(logdensity_def(LebesgueBase()^3, xf)) isa Float32
+    @test @inferred(logdensity_def(LebesgueBase(), x)) isa Float32
+    @test @inferred(logdensity_rel(StdNormal(), StdUniform(), x)) isa Float32
+    @test @inferred(logdensity_rel(StdNormal()^3, StdUniform()^3, uf)) isa Float32
+    @test @inferred(logdensity_rel(StdNormal()^3, StdUniform()^3, xf)) isa Float32
+    @test @inferred(logdensity_rel(Lebesgue(), Dirac(1f0), 2f0)) isa Float32
+    @test @inferred(logdensity_rel(Dirac(1f0), Lebesgue(), 1f0)) isa Float32
+    @test @inferred(logdensityof(Dirac(1f0), 1f0)) isa Float32
+    @test @inferred(logdensityof(weightedmeasure(0.5f0, StdNormal())^3, xf)) isa Float32
+    @test @inferred(logdensityof(weightedmeasure(static(0.5), StdNormal())^3, xf)) isa Float32
+    @test @inferred(logdensityof(weightedmeasure(0.5, StdNormal()), x)) isa Float32
+    @test @inferred(logdensityof(mix, x)) isa Float32
+    @test @inferred(logdensityof(mix^3, uf)) isa Float32
+    @test @inferred(logdensities(StdNormal(), randn(Float32, 4))) isa Vector{Float32}
+    @test @inferred(logdensities(StdNormal()^3, randn(Float32, 3, 4))) isa Vector{Float32}
+
+    # Weights that carry derivatives keep them:
+    @test ForwardDiff.derivative(w -> logdensityof(weightedmeasure(w, StdNormal()), 0.3), 0.1) ≈ 1
+    @test ForwardDiff.derivative(w -> logdensityof(weightedmeasure(w, StdNormal())^2, [0.3, 0.1]), 0.1) ≈ 2
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 114aef51..c95954a4 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -17,6 +17,7 @@ include("test_basics.jl")
 include("getdof.jl")
 include("shape_contract.jl")
 include("logdensities.jl")
+include("numtype.jl")
 include("transport.jl")
 include("smf.jl")
 include("domains.jl")
diff --git a/test/test_mooncake.jl b/test/test_mooncake.jl
index 95ce84b5..a6582c1d 100644
--- a/test/test_mooncake.jl
+++ b/test/test_mooncake.jl
@@ -10,7 +10,6 @@ using MeasureBase
 using MeasureBase: transport_to
 using MeasureBase: isneginf, isposinf, _adignore_call
 using MeasureBase: check_dof, require_insupport, _origin_depth
-using MeasureBase: logdensityof_rt
 
 _mooncake_gradient(f, x) = Mooncake.value_and_gradient!!(
     Mooncake.prepare_gradient_cache(f, x), f, x
@@ -27,7 +26,6 @@ _mooncake_gradient(f, x) = Mooncake.value_and_gradient!!(
         Mooncake.TestUtils.test_rule(rng, check_dof, StdNormal(), StdUniform(); is_primitive = true)
         Mooncake.TestUtils.test_rule(rng, require_insupport, StdNormal(), 0.5; is_primitive = true)
         Mooncake.TestUtils.test_rule(rng, _origin_depth, StdNormal(); is_primitive = true)
-        Mooncake.TestUtils.test_rule(rng, logdensityof_rt, StdNormal(), 0.5; is_primitive = true)
     end
 
     @testset "@_adignore is ignored" begin

From 87ebc03dc02b5dfa9ca9b5b19bbe07f1487bfa8f Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 17 Sep 2026 23:37:57 +0200
Subject: [PATCH 082/122] Add Reactant smoke tests and a traced standard normal
 quantile

The smoke tests live in their own project under test/reactant and are
not part of the default test suite. The Reactant extension provides the
standard normal quantile via erf_inv, since CHLO has no erfc_inv.

Created by generative AI.
---
 ext/MeasureBaseReactantExt.jl |  7 +++-
 test/reactant/Project.toml    | 12 ++++++
 test/reactant/runtests.jl     | 76 +++++++++++++++++++++++++++++++++++
 3 files changed, 94 insertions(+), 1 deletion(-)
 create mode 100644 test/reactant/Project.toml
 create mode 100644 test/reactant/runtests.jl

diff --git a/ext/MeasureBaseReactantExt.jl b/ext/MeasureBaseReactantExt.jl
index 53bc6b4b..c39be494 100644
--- a/ext/MeasureBaseReactantExt.jl
+++ b/ext/MeasureBaseReactantExt.jl
@@ -2,11 +2,16 @@
 
 module MeasureBaseReactantExt
 
-using Reactant: TracedRNumber
+using Reactant: Reactant, TracedRNumber
+using IrrationalConstants: sqrt2
 import MeasureBase
 using MeasureBase: RealValues, IntegerValues
 
 Base.in(::TracedRNumber{<:Real}, ::RealValues) = true
 Base.in(::TracedRNumber{<:Integer}, ::IntegerValues) = true
 
+# CHLO provides erf_inv but no erfc_inv, so the standard normal quantile
+# loses precision for arguments close to 0 and 1 in traced code:
+MeasureBase.Φinv(p::TracedRNumber) = Reactant.Ops.erf_inv(2 * p - 1) * sqrt2
+
 end # module MeasureBaseReactantExt
diff --git a/test/reactant/Project.toml b/test/reactant/Project.toml
new file mode 100644
index 00000000..775cd85e
--- /dev/null
+++ b/test/reactant/Project.toml
@@ -0,0 +1,12 @@
+[deps]
+ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+MeasureBase = "fa1605e6-acd5-459c-a1e6-7e635759db14"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+MeasureBase = {path = "../.."}
+
+[compat]
+Reactant = "0.2"
diff --git a/test/reactant/runtests.jl b/test/reactant/runtests.jl
new file mode 100644
index 00000000..f79e4754
--- /dev/null
+++ b/test/reactant/runtests.jl
@@ -0,0 +1,76 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+# Reactant smoke tests, not part of the default test suite. Run with
+# `julia --project=test/reactant test/reactant/runtests.jl` after
+# instantiating that project, or include this file in an environment that
+# provides Reactant.
+
+using Test
+using Reactant
+using MeasureBase
+using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic, Lebesgue, Dirac
+using MeasureBase: logdensities, logdensity_rel, weightedmeasure, superpose, restrict, mintegrate_exp
+using ArraysOfArrays: VectorOfSimilarVectors, sliced
+using Distributions: Normal, Exponential, Uniform, Beta
+
+Reactant.set_default_backend("cpu")
+
+# Compiles `f` for traced copies of `args` and compares with the plain result:
+function test_traced(f, args...; kwargs...)
+    expected = f(args...)
+    traced_args = map(Reactant.to_rarray, args)
+    result = @jit f(traced_args...)
+    @test _plain(result) ≈ _plain(expected) nans = true
+    return result
+end
+
+_plain(x::AbstractArray) = Array(x)
+_plain(x::Number) = Float64(x)
+
+@testset "Reactant" begin
+    x = randn(10)
+    X = randn(3, 20)
+
+    @testset "powers and batches" begin
+        test_traced(x -> logdensityof(StdNormal()^10, x), x)
+        test_traced(X -> logdensities(StdNormal(), X), X)
+        test_traced(X -> logdensities(StdNormal()^3, X), X)
+        test_traced(X -> logdensities(StdNormal()^3, sliced(X, 1)), X)
+        test_traced(X -> logdensities((StdNormal()^3)^4, reshape(X[:, 1:16], 3, 4, 4)), X)
+        test_traced(x -> logdensityof(StdUniform()^10, x), rand(10))
+        test_traced(x -> logdensityof(StdExponential()^10, x), rand(10))
+        test_traced(x -> logdensityof(weightedmeasure(0.3, StdNormal())^10, x), x)
+        test_traced(x -> logdensityof(Lebesgue()^10, x), x)
+    end
+
+    @testset "support masks" begin
+        xu = 2 .* rand(10) .- 0.5
+        test_traced(x -> logdensities(StdUniform(), x), xu)
+        test_traced(x -> logdensities(StdExponential(), x), xu)
+        test_traced(x -> logdensities(restrict(x -> x > 0, StdNormal()), x), xu)
+    end
+
+    @testset "relative densities" begin
+        xu = 2 .* rand(10) .- 0.5
+        test_traced(x -> logdensity_rel.(Ref(StdUniform()), Ref(StdExponential()), x), xu)
+        test_traced(x -> logdensity_rel.(Ref(StdNormal()), Ref(StdLogistic()), x), x)
+        test_traced(x -> logdensity_rel.(Ref(StdNormal()^10), Ref(StdLogistic()^10), Ref(x)), x)
+    end
+
+    @testset "superposition, density measures and spike mixtures" begin
+        mix = superpose(weightedmeasure(log(0.3), StdNormal()), weightedmeasure(log(0.7), StdLogistic()))
+        test_traced(x -> logdensities(mix, x), x)
+        test_traced(x -> logdensityof(mix^10, x), x)
+        dm = mintegrate_exp(x -> -abs(x), StdNormal())
+        test_traced(x -> logdensities(dm, x), x)
+        sm = SpikeMixture(StdNormal(), 0.2)
+        test_traced(x -> logdensities(sm, x), vcat(x, 0.0))
+    end
+
+    @testset "transport" begin
+        test_traced(x -> transport_to(StdUniform(), StdNormal()).(x), x)
+        test_traced(x -> transport_to(StdNormal(), StdUniform()).(x), rand(10))
+        test_traced(x -> transport_to(Normal(2, 3), StdNormal()).(x), x)
+        test_traced(x -> transport_to(StdNormal(), Exponential(2.0)).(x), rand(10))
+    end
+end

From da55d031ddad4f6edde5ea01114c48d6f3763e93 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 17 Sep 2026 23:37:57 +0200
Subject: [PATCH 083/122] Fix variate size contract and batched routing edge
 cases

Single variates of powers must have exactly the variate rank, batches of
scalar variates must be arrays of numbers, empty powers of measures with
unknown variate size evaluate to zero, checked_arg accepts variates in
flat storage, reshapes of nested variates report no flat storage,
array-valued sets distinguish element size from flat size, custom
standard measures rank lowest in promotion, distributions without an
array-like variate report an unknown size and product distributions
declare their preferred standard measure.

Created by generative AI.
---
 .../distribution_measure.jl                   |  2 ++
 ext/MeasureBaseDistributionsExt/product.jl    |  8 +++++
 src/combinators/power.jl                      | 15 ++++----
 src/combinators/reshape.jl                    | 13 ++++++-
 src/density-batched.jl                        | 35 +++++++++++++++----
 src/domains.jl                                |  5 ++-
 src/primitives/counting.jl                    |  2 +-
 src/standard/stdtraits.jl                     |  3 +-
 test/distributions/test_shape_contract.jl     |  7 ++++
 test/logdensities.jl                          | 20 +++++++++++
 test/shape_contract.jl                        | 19 ++++++++++
 11 files changed, 112 insertions(+), 17 deletions(-)

diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index 8d3eb8eb..4b6fb205 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -59,6 +59,8 @@ end
 
 @inline MeasureBase.massof(::DistributionMeasure) = static(1.0)
 
+@inline MeasureBase.mspace_elsize(d::Distribution) = MeasureBase.NoMSpaceElementSize{typeof(d)}()
+@inline MeasureBase.mspace_flatsize(d::Distribution) = MeasureBase.NoMSpaceElementSize{typeof(d)}()
 @inline MeasureBase.mspace_elsize(d::Distribution{Univariate}) = ()
 @inline MeasureBase.mspace_elsize(d::Distribution{<:ArrayLikeVariate}) = size(d)
 @inline MeasureBase.mspace_flatsize(d::Distribution{Univariate}) = ()
diff --git a/ext/MeasureBaseDistributionsExt/product.jl b/ext/MeasureBaseDistributionsExt/product.jl
index a050dc97..35817c4b 100644
--- a/ext/MeasureBaseDistributionsExt/product.jl
+++ b/ext/MeasureBaseDistributionsExt/product.jl
@@ -3,6 +3,10 @@
 @static if isdefined(Distributions, :Product)
     MeasureBase.AbstractMeasure(obj::Distributions.Product) = productmeasure(map(asmeasure, obj.v))
 
+    @inline function MeasureBase.preferred_stdmeasure(::Type{<:Distributions.Product{<:Any,T}}) where {T}
+        MeasureBase.preferred_stdmeasure(T)
+    end
+
     function AsMeasure{D}(::D) where {D<:Distributions.Product}
         throw(ArgumentError("Don't wrap Distributions.Product into MeasureBase.AsMeasure, use asmeasure to convert instead."))
     end
@@ -24,6 +28,10 @@ end
 @static if isdefined(Distributions, :ProductDistribution)
     MeasureBase.AbstractMeasure(obj::Distributions.ProductDistribution) = productmeasure(map(asmeasure, obj.dists))
 
+    @inline function MeasureBase.preferred_stdmeasure(::Type{<:Distributions.ProductDistribution{N,M,D}}) where {N,M,D<:AbstractArray}
+        MeasureBase.preferred_stdmeasure(eltype(D))
+    end
+
     function AsMeasure{D}(::D) where {D<:Distributions.ProductDistribution}
         throw(ArgumentError("Don't wrap Distributions.ProductDistribution into MeasureBase.AsMeasure, use asmeasure to convert instead."))
     end
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 590e576d..7eaf10f5 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -137,20 +137,21 @@ end
     static(0)
 end
 
+# Variates may be nested arrays of the power's shape or their flat storage:
 @propagate_inbounds function checked_arg(μ::PowerMeasure, x::AbstractArray{<:Any})
     @boundscheck begin
-        sz_μ = pwr_size(μ)
-        sz_x = size(x)
-        if sz_μ != sz_x
-            throw(ArgumentError("Size of variate doesn't match size of power measure"))
+        sz_x = maybestatic_size(x)
+        if sz_x != pwr_size(μ) && !_matches_flatsize(sz_x, mspace_flatsize(μ))
+            _throw_size_mismatch()
         end
     end
     return x
 end
 
-function checked_arg(μ::PowerMeasure, x::Any)
-    throw(ArgumentError("Size of variate doesn't match size of power measure"))
-end
+@inline _matches_flatsize(sz_x, sz_flat::SizeLike) = Tuple(sz_x) == Tuple(sz_flat)
+@inline _matches_flatsize(sz_x, ::NoMSpaceElementSize) = false
+
+checked_arg(μ::PowerMeasure, x::Any) = _throw_size_mismatch()
 
 massof(m::PowerMeasure) = massof(m.parent)^prod(m.axes)
 
diff --git a/src/combinators/reshape.jl b/src/combinators/reshape.jl
index 01f5be0d..ad83e472 100644
--- a/src/combinators/reshape.jl
+++ b/src/combinators/reshape.jl
@@ -56,4 +56,15 @@ mreshape(m::AbstractMeasure, sz::IntegerLike...) = mreshape(m, sz)
 mreshape(m::AbstractMeasure, sz::SizeLike) = pushfwd(Reshape(sz, some_mspace_elsize(m)), m)
 
 @inline mspace_elsize(μ::PushforwardMeasure{<:Reshape}) = μ.f.output_size
-@inline mspace_flatsize(μ::PushforwardMeasure{<:Reshape}) = μ.f.output_size
+# Reshaped variates have flat storage only if the reshaped elements are
+# numbers, the reshape of a nested variate has no flat form its density
+# kernel could consume:
+@inline function mspace_flatsize(μ::PushforwardMeasure{<:Reshape})
+    _reshaped_flatsize(mspace_flatsize(μ.origin), mspace_elsize(μ.origin), μ.f.output_size)
+end
+@inline function _reshaped_flatsize(sz_flat::SizeLike, sz_outer::SizeLike, sz_out)
+    _reshaped_flatsize(Val(length(_size_dims(sz_flat)) == length(_size_dims(sz_outer))), sz_out)
+end
+@inline _reshaped_flatsize(::Val{true}, sz_out) = sz_out
+@inline _reshaped_flatsize(::Val{false}, sz_out) = NoMSpaceElementSize{typeof(sz_out)}()
+@inline _reshaped_flatsize(::Any, ::Any, sz_out) = NoMSpaceElementSize{typeof(sz_out)}()
diff --git a/src/density-batched.jl b/src/density-batched.jl
index dd4910bf..05c91d69 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -67,6 +67,17 @@ struct NoFlatStorage end
     _batched_ld_flat(f, μ, X, _flat_storage(X), sz_flat)
 end
 
+# Batches of scalar variates are arrays of numbers, nested arrays are not
+# reinterpreted as flat storage:
+@inline function _batched_ld_sized(f::F, μ, X::AbstractArray, sz_flat::Tuple{}) where {F}
+    _batched_ld_flat(f, μ, X, _flat_scalar_storage(X), sz_flat)
+end
+
+@inline _flat_scalar_storage(X::AbstractArray{<:Number}) = X
+function _flat_scalar_storage(::AbstractArray)
+    throw(ArgumentError("A batch of scalar variates must be an array of numbers"))
+end
+
 @inline function _batched_ld_sized(f::F, μ, X::AbstractArray, ::NoMSpaceElementSize) where {F}
     map(x -> _pointwise_ld(f, μ, x), X)
 end
@@ -98,6 +109,9 @@ end
 
 @inline function _powered_ld_flat(f::F, μ::PowerMeasure, x, x_flat::AbstractArray, sz_flat) where {F}
     ν, n_pwr = _pwr_unwrap(μ)
+    if ndims(x_flat) != length(sz_flat)
+        _throw_size_mismatch()
+    end
     _check_flatsize(x_flat, sz_flat)
     _sum_leading_dims(_batched_kernel(f, ν, x_flat), n_pwr)
 end
@@ -107,16 +121,19 @@ end
 # Sum of the point-level densities over the elements of the variate:
 @inline function _powered_ld_pointwise(f::F, μ::PowerMeasure, x::AbstractArray) where {F}
     if maybestatic_size(x) != pwr_size(μ)
-        throw(ArgumentError("Size of variate doesn't match size of power measure"))
+        _throw_size_mismatch()
     end
-    sum(Base.Fix1(_pointwise_ld_dyn, (f, pwr_base(μ))), x)
+    init = zero(float(real_numtype(typeof(x))))
+    sum(Base.Fix1(_pointwise_ld_dyn, (f, pwr_base(μ))), x; init = init)
 end
 
+@noinline _throw_size_mismatch() = throw(ArgumentError("Size of variate doesn't match size of measure"))
+
 function _powered_ld_pointwise(f::F, ::PowerMeasure, x) where {F}
     throw(ArgumentError("Variate of a power measure must be an array"))
 end
 
-@inline _pointwise_ld_dyn((f, μ), x) = dynamic(_pointwise_ld(f, μ, x))
+@inline _pointwise_ld_dyn((f, μ), x) = _dynamic_logd(_pointwise_ld(f, μ, x), x)
 
 @inline _pwr_unwrap(μ) = (μ, static(0))
 @inline function _pwr_unwrap(μ::PowerMeasure)
@@ -141,7 +158,7 @@ end
 @inline function _check_flatsize(A::AbstractArray, sz_flat::SizeLike)
     n = length(sz_flat)
     if ndims(A) < n || ntuple(i -> size(A, i), Val(n)) != Tuple(sz_flat)
-        throw(ArgumentError("Size of variate doesn't match size of measure"))
+        _throw_size_mismatch()
     end
     return nothing
 end
@@ -153,7 +170,7 @@ end
 
 # Powers of primitive measures have log-density zero relative to their base:
 @inline function _batched_kernel(::typeof(logdensity_def), ν::PrimitiveMeasure, A::AbstractArray)
-    FillArrays.Zeros{Float64}(size(A))
+    FillArrays.Zeros{_logd_numtype(A)}(size(A))
 end
 
 @inline _batched_ld_generic(f::F, ν, A::AbstractArray) where {F} = _batched_ld_byflatsize(f, ν, A, mspace_flatsize(ν))
@@ -163,8 +180,14 @@ end
 # Static results of point kernels are made dynamic, to keep reductions
 # over them type stable.
 @inline function _batched_ld_byflatsize(f::F, ν, A::AbstractArray, ::Tuple{}) where {F}
-    Broadcast.instantiate(Broadcast.broadcasted(dynamic ∘ Base.Fix1(f, ν), A))
+    Broadcast.instantiate(Broadcast.broadcasted(_DynamicLogd(f, ν), A))
+end
+
+struct _DynamicLogd{F,M} <: Function
+    f::F
+    ν::M
 end
+@inline (k::_DynamicLogd)(x) = _dynamic_logd(k.f(k.ν, x), x)
 
 # Array variates: map over the variate slices, or evaluate directly if `A`
 # is a single variate.
diff --git a/src/domains.jl b/src/domains.jl
index a6a4e800..643f1ec0 100644
--- a/src/domains.jl
+++ b/src/domains.jl
@@ -431,5 +431,8 @@ _combinesets_cat(f_c, α::SetLike, β::SetLike) = CombinedSet(f_c, α, β)
 @inline _valueset_elsize(::IntegerValues) = ()
 @inline _valueset_elsize(::BoundedInts) = ()
 @inline _valueset_elsize(::IntervalSets.AbstractInterval) = ()
-@inline _valueset_elsize(s::CartesianPower) = _cat_sizes(_valueset_elsize(pwr_base(s)), pwr_size(s))
+@inline _valueset_elsize(s::CartesianPower) = pwr_size(s)
 @inline _valueset_elsize(s) = NoMSpaceElementSize{typeof(s)}()
+
+@inline _valueset_flatsize(s::CartesianPower) = _cat_sizes(_valueset_flatsize(pwr_base(s)), pwr_size(s))
+@inline _valueset_flatsize(s) = _valueset_elsize(s)
diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl
index 0586ba2e..8bcdb052 100644
--- a/src/primitives/counting.jl
+++ b/src/primitives/counting.jl
@@ -29,7 +29,7 @@ basemeasure(::Counting) = CountingBase()
 Counting() = Counting(ℤ)
 
 @inline mspace_elsize(μ::Counting) = _valueset_elsize(μ.support)
-@inline mspace_flatsize(μ::Counting) = _valueset_elsize(μ.support)
+@inline mspace_flatsize(μ::Counting) = _valueset_flatsize(μ.support)
 
 testvalue(::Type{T}, d::Counting) where {T} = testvalue(T, d.support)
 
diff --git a/src/standard/stdtraits.jl b/src/standard/stdtraits.jl
index 69752356..c747f21e 100644
--- a/src/standard/stdtraits.jl
+++ b/src/standard/stdtraits.jl
@@ -53,7 +53,7 @@ to `StdNormal`.
 function promote_stdmeasure end
 
 @inline function promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:StdMeasure,B<:StdMeasure}
-    ifelse(_stdmeasure_rank(A) >= _stdmeasure_rank(B), A, B)
+    _stdmeasure_rank(A) >= _stdmeasure_rank(B) ? A : B
 end
 
 @inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{B}) where {B} = B
@@ -70,6 +70,7 @@ end
     promote_stdmeasure(promote_stdmeasure(A, B), Cs...)
 end
 
+@inline _stdmeasure_rank(::Type{<:StdMeasure}) = 0
 @inline _stdmeasure_rank(::Type{StdUniform}) = 1
 @inline _stdmeasure_rank(::Type{StdExponential}) = 2
 @inline _stdmeasure_rank(::Type{StdLogistic}) = 3
diff --git a/test/distributions/test_shape_contract.jl b/test/distributions/test_shape_contract.jl
index d02c5680..b4f11ad7 100644
--- a/test/distributions/test_shape_contract.jl
+++ b/test/distributions/test_shape_contract.jl
@@ -31,4 +31,11 @@ using LinearAlgebra: I
     @test @inferred(preferred_stdmeasure(productmeasure((a = Dirac(1.0), b = asmeasure(Beta(2, 3)))))) === StdUniform
     @test @inferred(preferred_stdmeasure(productmeasure((a = asmeasure(Poisson(2)), b = asmeasure(Beta(2, 3)))))) <: NoStdTransport
     @test @inferred(preferred_stdmeasure(productmeasure([asmeasure(Normal(i, 1)) for i in 1:3]))) === StdNormal
+    @test @inferred(preferred_stdmeasure(product_distribution([Beta(2, 3), Beta(1, 1)]))) === StdUniform
+
+    lkj = asmeasure(LKJCholesky(3, 1.0))
+    @test @inferred(mspace_elsize(lkj)) isa MeasureBase.NoMSpaceElementSize
+    @test @inferred(mspace_flatsize(lkj)) isa MeasureBase.NoMSpaceElementSize
+    X_lkj = [rand(LKJCholesky(3, 1.0)) for _ in 1:3]
+    @test logdensities(lkj, X_lkj) ≈ logdensityof.(Ref(lkj), X_lkj)
 end
diff --git a/test/logdensities.jl b/test/logdensities.jl
index e823a138..25e31328 100644
--- a/test/logdensities.jl
+++ b/test/logdensities.jl
@@ -60,9 +60,29 @@ stdnormal_ld(x) = -(x^2 + log2π) / 2
         @test logdensities(mix, X) ≈ logdensityof.(Ref(mix), X)
         @test logdensityof(mix^4, X[:, 1]) ≈ sum(logdensityof.(Ref(mix), X[:, 1]))
         @test logdensities(mix^4, sliced(X, 1)) ≈ vec(sum(logdensityof.(Ref(mix), X), dims = 1))
+        @test @inferred(logdensityof(mix^0, Float64[])) == 0
+        @test logdensities(mix^0, [Float64[], Float64[]]) == [0.0, 0.0]
+    end
+
+    @testset "flat and nested variate forms agree" begin
+        for (μ, x_flat) in (
+            ((StdNormal()^3)^2, randn(3, 2)),
+            ((StdNormal()^(2, 3))^4, randn(2, 3, 4)),
+        )
+            x_nested = sliced(x_flat, length(MeasureBase.mspace_flatsize(MeasureBase.pwr_base(μ))))
+            @test logdensityof(μ, x_flat) ≈ logdensityof(μ, x_nested)
+            @test MeasureBase.checked_arg(μ, x_flat) === x_flat
+            @test MeasureBase.checked_arg(μ, x_nested) === x_nested
+            @test_throws ArgumentError MeasureBase.checked_arg(μ, randn(7))
+        end
     end
 
     @testset "size mismatch" begin
+        @test_throws ArgumentError logdensityof(StdNormal()^3, randn(3, 4))
+        @test_throws ArgumentError logdensityof(StdNormal()^(2, 3), randn(2, 3, 1))
+        @test_throws ArgumentError logdensityof(StdNormal()^3, randn(4))
+        @test_throws ArgumentError logdensityof(StdNormal()^3, 1.0)
+        @test_throws ArgumentError logdensities(StdNormal(), VectorOfSimilarVectors(randn(3, 5)))
         @test_throws ArgumentError logdensities(StdNormal()^3, [randn(3), randn(2)])
         @test_throws ArgumentError logdensities(
             StdNormal()^3,
diff --git a/test/shape_contract.jl b/test/shape_contract.jl
index 31920c9f..779b40b2 100644
--- a/test/shape_contract.jl
+++ b/test/shape_contract.jl
@@ -9,6 +9,13 @@ using MeasureBase: mreshape, productmeasure, weightedmeasure, pushfwd, mbind, re
 using IntervalSets: (..)
 using StaticArrays: SVector, Size
 using Static: static
+using MeasureBase: size2length
+using MeasureBase: setcartpower, ℝ, testvalue
+
+_flat_iter(x::Number) = (x,)
+_flat_iter(x::AbstractArray) = Iterators.flatten(map(_flat_iter, x))
+
+struct _CustomStd <: MeasureBase.StdMeasure end
 
 @testset "shape contract" begin
     @testset "mspace_elsize and mspace_flatsize" begin
@@ -39,6 +46,15 @@ using Static: static
         @test @inferred(mspace_elsize(restrict(x -> x > 0, StdNormal()))) === ()
         @test @inferred(mspace_elsize(mreshape(StdNormal()^6, (2, 3)))) == (2, 3)
         @test @inferred(mspace_flatsize(mreshape(StdNormal()^6, (2, 3)))) == (2, 3)
+        @test @inferred(mspace_elsize(mreshape((StdNormal()^2)^6, (2, 3)))) == (2, 3)
+        @test @inferred(mspace_flatsize(mreshape((StdNormal()^2)^6, (2, 3)))) isa NoMSpaceElementSize
+
+        s = setcartpower(setcartpower(ℝ, 2), 3)
+        @test @inferred(mspace_elsize(Lebesgue(s))) == (3,)
+        @test @inferred(mspace_flatsize(Lebesgue(s))) == (2, 3)
+        for μ in (StdNormal(), StdNormal()^3, (StdNormal()^2)^3, Dirac([1.0, 2.0]), Dirac(3.0))
+            @test size2length(mspace_flatsize(μ)) == length(vec(collect(Iterators.flatten(_flat_iter(testvalue(μ))))))
+        end
 
         @test @inferred(mspace_elsize(productmeasure((a = StdNormal(), b = StdUniform())))) isa NoMSpaceElementSize
         @test @inferred(mspace_flatsize(mbind(x -> StdNormal()^2, StdUniform()))) isa NoMSpaceElementSize
@@ -81,5 +97,8 @@ using Static: static
         @test @inferred(promote_stdmeasure(NoStdTransport{Int}, AnyStdMeasure)) === NoStdTransport{Int}
         @test @inferred(promote_stdmeasure(AnyStdMeasure, NoStdTransport{Int})) === NoStdTransport{Int}
         @test @inferred(promote_stdmeasure(StdUniform, StdExponential, AnyStdMeasure, StdLogistic)) === StdLogistic
+        @test @inferred(promote_stdmeasure(_CustomStd, StdUniform)) === StdUniform
+        @test @inferred(promote_stdmeasure(_CustomStd, AnyStdMeasure)) === _CustomStd
+        @test @inferred(preferred_stdmeasure(_CustomStd())) === _CustomStd
     end
 end

From f726c11c6ab3234a84a35616521955caf806a8e2 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 17 Sep 2026 23:53:25 +0200
Subject: [PATCH 084/122] Route measures of unknown variate size through the
 batched kernel

Batches of points of measures without a known flat variate size reach
batched_logdensityof_impl as arrays of points, with a lazy broadcast of
the point kernel as the default, so batched kernels for structured and
ragged variates can be provided. Superpositions, spike mixtures, density
measures, restricted and half measures declare the flat variate size of
their components at the type level, so powers of scalar superpositions
take the fused path. Support checks of powers run over the flat variate
storage. Restricted and half measures declare no standard-measure
transport, arrays of marginals of mixed types combine their transport
preferences at run time.

Created by generative AI.
---
 .../distribution_measure.jl                   |  2 ++
 src/combinators/half.jl                       |  3 +-
 src/combinators/power.jl                      | 34 +++++++++++--------
 src/combinators/restricted.jl                 |  1 +
 src/combinators/spikemixture.jl               |  4 +++
 src/combinators/superpose.jl                  | 11 ++++++
 src/combinators/weighted.jl                   |  1 +
 src/density-batched.jl                        | 20 +++++++----
 src/density.jl                                |  4 +++
 src/mspace.jl                                 | 17 ++++++++++
 src/primitives/counting.jl                    |  3 ++
 src/primitives/dirac.jl                       |  1 +
 src/primitives/lebesgue.jl                    |  3 ++
 src/standard/stdmeasure.jl                    |  1 +
 src/standard/stdtraits.jl                     | 13 ++++++-
 test/shape_contract.jl                        |  4 ++-
 16 files changed, 97 insertions(+), 25 deletions(-)

diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index 4b6fb205..8a99df31 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -67,6 +67,8 @@ end
 @inline MeasureBase.mspace_flatsize(d::Distribution{<:ArrayLikeVariate}) = size(d)
 @inline MeasureBase.mspace_elsize(m::DistributionMeasure) = MeasureBase.mspace_elsize(m.obj)
 @inline MeasureBase.mspace_flatsize(m::DistributionMeasure) = MeasureBase.mspace_flatsize(m.obj)
+@inline MeasureBase.mspace_flatsize(::Type{<:Distribution{Univariate}}) = ()
+@inline MeasureBase.mspace_flatsize(::Type{AsMeasure{D}}) where {D<:Distribution} = MeasureBase.mspace_flatsize(D)
 
 @inline MeasureBase.preferred_stdmeasure(::Type{AsMeasure{D}}) where {D<:Distribution} = MeasureBase.preferred_stdmeasure(D)
 
diff --git a/src/combinators/half.jl b/src/combinators/half.jl
index 906978ca..e4c35788 100644
--- a/src/combinators/half.jl
+++ b/src/combinators/half.jl
@@ -6,7 +6,8 @@ end
 
 @inline mspace_elsize(μ::Half) = mspace_elsize(μ.parent)
 @inline mspace_flatsize(μ::Half) = mspace_flatsize(μ.parent)
-@inline preferred_stdmeasure(::Type{<:Half{M}}) where {M} = preferred_stdmeasure(M)
+@inline mspace_flatsize(::Type{<:Half{M}}) where {M} = mspace_flatsize(M)
+@inline preferred_stdmeasure(::Type{MU}) where {MU<:Half} = NoStdTransport{MU}
 
 function Base.show(io::IO, μ::Half)
     print(io, "Half")
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 7eaf10f5..3a2f77aa 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -106,24 +106,28 @@ end
 @inline logdensityof_impl(μ::PowerMeasure, x) = _powered_ld(logdensityof_impl, μ, x)
 @inline logdensity_def(μ::PowerMeasure, x) = _powered_ld(logdensity_def, μ, x)
 
-@inline function insupport(μ::PowerMeasure, x)
-    p = μ.parent
-    all(x) do xj
-        # https://github.com/SciML/Static.jl/issues/36
-        dynamic(insupport(p, xj))
-    end
+# Support checks of powers run over the flat variate storage where the base
+# measure has scalar variates, elementwise otherwise:
+@inline function insupport(μ::PowerMeasure, x::AbstractArray)
+    _powered_insupport(μ, x, _flat_storage(x), mspace_flatsize(μ))
+end
+
+@inline function _powered_insupport(μ::PowerMeasure, x, x_flat::AbstractArray, ::SizeLike)
+    ν, _ = _pwr_unwrap(μ)
+    _powered_insupport_flat(ν, x_flat, mspace_flatsize(ν))
 end
+@inline function _powered_insupport_flat(ν, x_flat::AbstractArray, ::Tuple{})
+    _all_insupport(broadcast(_insupport_bool ∘ Base.Fix1(insupport, ν), x_flat))
+end
+@inline _powered_insupport_flat(ν, x_flat::AbstractArray, ::Any) = _powered_insupport_elementwise(ν, x_flat)
+@inline _powered_insupport(μ::PowerMeasure, x, ::Any, ::Any) = _powered_insupport_elementwise(pwr_base(μ), x)
 
-_all(A) = all(A)
-_all(::AbstractArray{NoFastInsupport{T}}) where {T} = NoFastInsupport{T}()
+@inline function _powered_insupport_elementwise(ν, x::AbstractArray)
+    _all_insupport(broadcast(_insupport_bool ∘ Base.Fix1(insupport, ν), x))
+end
 
-@inline function insupport(μ::PowerMeasure, x::AbstractArray)
-    p = μ.parent
-    insupp = broadcast(x) do xj
-        # https://github.com/SciML/Static.jl/issues/36
-        dynamic(insupport(p, xj))
-    end
-    _all(insupp)
+function insupport(μ::PowerMeasure, x)
+    mapreduce(_insupport_bool ∘ Base.Fix1(insupport, pwr_base(μ)), _insupport_and, x)
 end
 
 @inline getdof(μ::PowerMeasure) = getdof(μ.parent) * size2length(axes2size(μ.axes))
diff --git a/src/combinators/restricted.jl b/src/combinators/restricted.jl
index 792d3a5b..f80e4dc4 100644
--- a/src/combinators/restricted.jl
+++ b/src/combinators/restricted.jl
@@ -5,6 +5,7 @@ end
 
 @inline mspace_elsize(μ::RestrictedMeasure) = mspace_elsize(μ.base)
 @inline mspace_flatsize(μ::RestrictedMeasure) = mspace_flatsize(μ.base)
+@inline mspace_flatsize(::Type{<:RestrictedMeasure{<:Any,M}}) where {M} = mspace_flatsize(M)
 
 @inline logdensity_def(d::RestrictedMeasure, x) = logdensity_def(d.base, x)
 
diff --git a/src/combinators/spikemixture.jl b/src/combinators/spikemixture.jl
index 216d66f9..64033840 100644
--- a/src/combinators/spikemixture.jl
+++ b/src/combinators/spikemixture.jl
@@ -40,3 +40,7 @@ end
 testvalue(::Type{T}, μ::SpikeMixture) where {T} = zero(T)
 
 insupport(μ::SpikeMixture, x) = _insupport_mask(insupport(μ.m, x)) | iszero(x)
+
+
+@inline mspace_flatsize(μ::SpikeMixture) = _scalar_or_unknown(mspace_flatsize(μ.m))
+@inline mspace_flatsize(::Type{<:SpikeMixture{M}}) where {M} = _scalar_or_unknown(mspace_flatsize(M))
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index 369cbd5d..cb8073fe 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -144,3 +144,14 @@ end
 @inline function insupport(d::SuperpositionMeasure, x)
     mapreduce(c -> _insupport_mask(insupport(c, x)), |, values(d.components))
 end
+
+
+@inline mspace_flatsize(μ::SuperpositionMeasure) = mspace_flatsize(typeof(μ))
+@inline mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:AbstractArray} = _scalar_or_unknown(mspace_flatsize(eltype(C)))
+@inline mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:Tuple} = _common_scalar_flatsize(C)
+@generated function _common_scalar_flatsize(::Type{C}) where {C<:Tuple}
+    args = [:(mspace_flatsize($T)) for T in C.parameters]
+    :(_all_scalar_sizes($(args...)))
+end
+@inline _all_scalar_sizes(::Tuple{}...) = ()
+@inline _all_scalar_sizes(szs...) = NoMSpaceElementSize{typeof(szs)}()
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index aa3feb3d..6341e654 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -41,6 +41,7 @@ end
 
 @inline mspace_elsize(μ::WeightedMeasure) = mspace_elsize(μ.base)
 @inline mspace_flatsize(μ::WeightedMeasure) = mspace_flatsize(μ.base)
+@inline mspace_flatsize(::Type{<:WeightedMeasure{<:Any,M}}) where {M} = mspace_flatsize(M)
 
 massof(w::WeightedMeasure) = exp(w.logweight) * massof(w.base)
 
diff --git a/src/density-batched.jl b/src/density-batched.jl
index 05c91d69..03f7a138 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -78,9 +78,14 @@ function _flat_scalar_storage(::AbstractArray)
     throw(ArgumentError("A batch of scalar variates must be an array of numbers"))
 end
 
+# Measures of unknown variate size take `X` as an array of points:
 @inline function _batched_ld_sized(f::F, μ, X::AbstractArray, ::NoMSpaceElementSize) where {F}
-    map(x -> _pointwise_ld(f, μ, x), X)
+    _batched_kernel(f, μ, X)
+end
+@inline function _batched_ld_sized(f::F, μ::PowerMeasure, X::AbstractArray, ::NoMSpaceElementSize) where {F}
+    Broadcast.instantiate(Broadcast.broadcasted(_DynamicLogd(_pointwise_ld_fixed(f), μ), X))
 end
+@inline _pointwise_ld_fixed(f::F) where {F} = (μ, x) -> _pointwise_ld(f, μ, x)
 
 @inline function _batched_ld_flat(f::F, μ, X, X_flat::AbstractArray, sz_flat) where {F}
     ν, n_pwr = _pwr_unwrap(μ)
@@ -118,19 +123,20 @@ end
 
 @inline _powered_ld_flat(f::F, μ::PowerMeasure, x, ::NoFlatStorage, sz_flat) where {F} = _powered_ld_pointwise(f, μ, x)
 
-# Sum of the point-level densities over the elements of the variate:
+# Sum of the point-level densities of the base measure over the elements
+# of the variate, evaluated via the batched kernel of the base measure:
 @inline function _powered_ld_pointwise(f::F, μ::PowerMeasure, x::AbstractArray) where {F}
     if maybestatic_size(x) != pwr_size(μ)
         _throw_size_mismatch()
     end
-    init = zero(float(real_numtype(typeof(x))))
-    sum(Base.Fix1(_pointwise_ld_dyn, (f, pwr_base(μ))), x; init = init)
+    ν = pwr_base(μ)
+    _sum_leading_dims(_batched_ld_sized(f, ν, x, mspace_flatsize(ν)), static(ndims(x)))
 end
 
 @noinline _throw_size_mismatch() = throw(ArgumentError("Size of variate doesn't match size of measure"))
 
-function _powered_ld_pointwise(f::F, ::PowerMeasure, x) where {F}
-    throw(ArgumentError("Variate of a power measure must be an array"))
+function _powered_ld_pointwise(f::F, μ::PowerMeasure, x) where {F}
+    throw(ArgumentError("Variates of powers of measures must be arrays, and flat variate storage requires a base measure of known variate size"))
 end
 
 @inline _pointwise_ld_dyn((f, μ), x) = _dynamic_logd(_pointwise_ld(f, μ, x), x)
@@ -197,7 +203,7 @@ end
 
 # Variates of unknown size: the elements of `A` are the variates.
 @inline function _batched_ld_byflatsize(f::F, ν, A::AbstractArray, ::NoMSpaceElementSize) where {F}
-    map(Base.Fix1(f, ν), A)
+    Broadcast.instantiate(Broadcast.broadcasted(_DynamicLogd(f, ν), A))
 end
 
 @inline _batched_ld_slices(f::F, ν, A::AbstractArray{<:Any,N}, ::Val{N}) where {F,N} = f(ν, A)
diff --git a/src/density.jl b/src/density.jl
index 900332ae..7c56f43c 100644
--- a/src/density.jl
+++ b/src/density.jl
@@ -220,6 +220,10 @@ mintegrate_exp(log_f, μ::AbstractMeasure) = DensityMeasure(as_integrand_exp(log
 
 basemeasure(μ::DensityMeasure) = μ.base
 
+@inline mspace_elsize(μ::DensityMeasure) = mspace_elsize(μ.base)
+@inline mspace_flatsize(μ::DensityMeasure) = mspace_flatsize(μ.base)
+@inline mspace_flatsize(::Type{<:DensityMeasure{<:Any,B}}) where {B} = mspace_flatsize(B)
+
 logdensity_def(μ::DensityMeasure, x) = logdensityof(μ.f, x)
 
 density_def(μ::DensityMeasure, x) = densityof(μ.f, x)
diff --git a/src/mspace.jl b/src/mspace.jl
index f7fe706c..ca0d269e 100644
--- a/src/mspace.jl
+++ b/src/mspace.jl
@@ -83,3 +83,20 @@ _mspace_some_elsize_impl(μ::AbstractMeasure, ::NoMSpaceElementSize) =
 @inline _value_flatsize(::Number) = ()
 @inline _value_flatsize(x::AbstractArray{<:Number}) = maybestatic_size(x)
 @inline _value_flatsize(x) = NoMSpaceElementSize{typeof(x)}()
+
+@inline _scalar_or_unknown(::Tuple{}) = ()
+@inline _scalar_or_unknown(sz::NoMSpaceElementSize) = sz
+@inline _scalar_or_unknown(sz) = NoMSpaceElementSize{typeof(sz)}()
+
+
+"""
+    MeasureBase.mspace_flatsize(::Type{MU})
+
+The flat variate size of measures of type `MU`, if it is determined by the
+type alone, e.g. `()` for measures with scalar variates. Returns
+`NoMSpaceElementSize{MU}()` otherwise.
+
+Composite measures use it to determine the flat size of their variates
+without inspecting each component.
+"""
+@inline mspace_flatsize(::Type{MU}) where {MU} = NoMSpaceElementSize{MU}()
diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl
index 8bcdb052..f4335526 100644
--- a/src/primitives/counting.jl
+++ b/src/primitives/counting.jl
@@ -6,6 +6,7 @@ struct CountingBase <: PrimitiveMeasure end
 
 @inline mspace_elsize(::CountingBase) = ()
 @inline mspace_flatsize(::CountingBase) = ()
+@inline mspace_flatsize(::Type{CountingBase}) = ()
 
 insupport(::CountingBase, x) = true
 
@@ -30,6 +31,8 @@ Counting() = Counting(ℤ)
 
 @inline mspace_elsize(μ::Counting) = _valueset_elsize(μ.support)
 @inline mspace_flatsize(μ::Counting) = _valueset_flatsize(μ.support)
+@inline mspace_flatsize(::Type{<:Counting{IntegerValues}}) = ()
+@inline mspace_flatsize(::Type{<:Counting{<:BoundedInts}}) = ()
 
 testvalue(::Type{T}, d::Counting) where {T} = testvalue(T, d.support)
 
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 3b134463..457d2ce1 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -45,6 +45,7 @@ insupport(d::Dirac, x) = x == d.x
 
 @inline mspace_elsize(μ::Dirac) = _value_elsize(μ.x)
 @inline mspace_flatsize(μ::Dirac) = _value_flatsize(μ.x)
+@inline mspace_flatsize(::Type{<:Dirac{<:Number}}) = ()
 
 @propagate_inbounds function checked_arg(μ::Dirac, x)
     @boundscheck insupport(μ, x) || throw(ArgumentError("Invalid variate for measure"))
diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl
index 7bd718cf..1604b112 100644
--- a/src/primitives/lebesgue.jl
+++ b/src/primitives/lebesgue.jl
@@ -6,6 +6,7 @@ struct LebesgueBase <: PrimitiveMeasure end
 
 @inline mspace_elsize(::LebesgueBase) = ()
 @inline mspace_flatsize(::LebesgueBase) = ()
+@inline mspace_flatsize(::Type{LebesgueBase}) = ()
 
 massof(::LebesgueBase, s::Interval) = width(s)
 
@@ -53,6 +54,8 @@ Lebesgue() = Lebesgue(ℝ)
 
 @inline mspace_elsize(μ::Lebesgue) = _valueset_elsize(μ.support)
 @inline mspace_flatsize(μ::Lebesgue) = _valueset_flatsize(μ.support)
+@inline mspace_flatsize(::Type{<:Lebesgue{RealValues}}) = ()
+@inline mspace_flatsize(::Type{<:Lebesgue{<:IntervalSets.AbstractInterval}}) = ()
 
 testvalue(::Type{T}, d::Lebesgue) where {T} = testvalue(T, d.support)::T
 
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index 367e5891..bb88b0ae 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -13,6 +13,7 @@ const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}}
 
 @inline mspace_elsize(::StdMeasure) = ()
 @inline mspace_flatsize(::StdMeasure) = ()
+@inline mspace_flatsize(::Type{<:StdMeasure}) = ()
 
 @inline check_dof(::StdMeasure, ::StdMeasure) = nothing
 
diff --git a/src/standard/stdtraits.jl b/src/standard/stdtraits.jl
index c747f21e..1562c3a5 100644
--- a/src/standard/stdtraits.jl
+++ b/src/standard/stdtraits.jl
@@ -79,12 +79,23 @@ end
 
 @inline preferred_stdmeasure(::Type{<:PowerMeasure{M}}) where {M} = preferred_stdmeasure(M)
 @inline preferred_stdmeasure(::Type{<:WeightedMeasure{<:Any,M}}) where {M} = preferred_stdmeasure(M)
-@inline preferred_stdmeasure(::Type{<:RestrictedMeasure{<:Any,M}}) where {M} = preferred_stdmeasure(M)
+# Transports of the base don't transport the restricted measure:
+@inline preferred_stdmeasure(::Type{MU}) where {MU<:RestrictedMeasure} = NoStdTransport{MU}
 @inline preferred_stdmeasure(::Type{<:PushforwardMeasure{<:Any,<:Any,M}}) where {M} = preferred_stdmeasure(M)
 @inline preferred_stdmeasure(::Type{<:Dirac}) = AnyStdMeasure
 
 @inline preferred_stdmeasure(::Type{<:ProductMeasure{M}}) where {M<:AbstractArray} = preferred_stdmeasure(eltype(M))
 
+# Arrays of marginals of mixed types combine their preferences at run time:
+@inline function preferred_stdmeasure(μ::ProductMeasure{<:AbstractArray})
+    _array_product_stdmeasure(μ, preferred_stdmeasure(typeof(μ)))
+end
+@inline _array_product_stdmeasure(μ, S::Type) = S
+function _array_product_stdmeasure(μ::ProductMeasure{<:AbstractArray{M}}, ::Type{NoStdTransport{MU}}) where {M,MU}
+    isconcretetype(M) && return NoStdTransport{MU}
+    mapreduce(preferred_stdmeasure, promote_stdmeasure, marginals(μ); init = AnyStdMeasure)
+end
+
 @inline function preferred_stdmeasure(::Type{<:ProductMeasure{M}}) where {M<:Tuple}
     _promote_stdmeasure_oftypes(M)
 end
diff --git a/test/shape_contract.jl b/test/shape_contract.jl
index 779b40b2..7c1ebc20 100644
--- a/test/shape_contract.jl
+++ b/test/shape_contract.jl
@@ -66,8 +66,10 @@ struct _CustomStd <: MeasureBase.StdMeasure end
             @test @inferred(preferred_stdmeasure(S()^3)) === S
             @test @inferred(preferred_stdmeasure(weightedmeasure(0.1, S()))) === S
             @test @inferred(preferred_stdmeasure(pushfwd(exp, S()))) === S
-            @test @inferred(preferred_stdmeasure(restrict(x -> x > 0, S()))) === S
         end
+        # Transports of the base measure don't transport restricted measures:
+        @test @inferred(preferred_stdmeasure(restrict(x -> x > 0, StdNormal()))) <: NoStdTransport
+        @test @inferred(preferred_stdmeasure(MeasureBase.Half(StdNormal()))) <: NoStdTransport
 
         @test @inferred(preferred_stdmeasure(Dirac(2.0))) === AnyStdMeasure
         @test @inferred(preferred_stdmeasure(Lebesgue())) <: NoStdTransport

From 02166ce6e716a53f914be56d1fc20d9532015e84 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 00:09:42 +0200
Subject: [PATCH 085/122] Add structural batched density kernels and batched
 stream consumption

Products over arrays of marginals, weighted measures and vcat-combined
measures now implement batched_logdensityof_impl, so batches of their
flat variates are evaluated in fused operations. Combined measures use
the new batched_logdensityof_with_rest, the batched form of the stream
consumption protocol. Powers that reach the batched extension point
route back into the batched core, and mcombine only merges products
with marginals of equal type.

Created by generative AI.
---
 src/combinators/combined.jl | 27 +++++++++++++-
 src/combinators/power.jl    |  1 +
 src/combinators/product.jl  | 21 +++++++++++
 src/combinators/weighted.jl |  4 +++
 src/density-batched.jl      | 56 +++++++++++++++++++++++++++--
 src/density-core.jl         |  6 +++-
 test/logdensities.jl        | 71 ++++++++++++++++++++++++++++++++++++-
 7 files changed, 180 insertions(+), 6 deletions(-)

diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index 3d4d3192..9ebe980b 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -84,7 +84,7 @@ end
     _mcombine_product_shortcut(f_c, marginals(α), marginals(β), α, β)
 end
 
-_mcombine_product_shortcut(::typeof(vcat), ma::AbstractVector, mb::AbstractVector, α, β) =
+_mcombine_product_shortcut(::typeof(vcat), ma::AbstractVector{T}, mb::AbstractVector{T}, α, β) where {T} =
     productmeasure(vcat(ma, mb))
 _mcombine_product_shortcut(::typeof(merge), ma::NamedTuple, mb::NamedTuple, α, β) =
     productmeasure(merge(ma, mb))
@@ -125,6 +125,15 @@ end
 
 @inline insupport(μ::CombinedMeasure, ab) = NoFastInsupport{typeof(μ)}()
 
+@inline function mspace_flatsize(μ::CombinedMeasure{typeof(vcat)})
+    _vcat_flatsize(mspace_flatsize(μ.α), mspace_flatsize(μ.β))
+end
+
+@inline _vcat_flatsize(a::SizeLike, b::SizeLike) = (size2length(a) + size2length(b),)
+@inline _vcat_flatsize(a::NoMSpaceElementSize, ::Any) = a
+@inline _vcat_flatsize(::Any, b::NoMSpaceElementSize) = b
+@inline _vcat_flatsize(a::NoMSpaceElementSize, ::NoMSpaceElementSize) = a
+
 @inline getdof(μ::CombinedMeasure) = getdof(μ.α) + getdof(μ.β)
 @inline fast_dof(μ::CombinedMeasure) = fast_dof(μ.α) + fast_dof(μ.β)
 
@@ -176,6 +185,22 @@ function _combined_ld_impl(f_c, μ::CombinedMeasure, ab)
     return logdensityof(tpm_α, a) + logdensityof(μ.β, b)
 end
 
+function batched_logdensityof_impl(μ::CombinedMeasure{typeof(vcat)}, A::AbstractArray)
+    ℓ, _, A_rest = batched_logdensityof_with_rest(μ, A)
+    if size(A_rest, 1) != 0
+        throw(ArgumentError("Variate streams too long during batched density evaluation of a combined measure"))
+    end
+    return ℓ
+end
+
+function batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, A::AbstractArray)
+    ℓ_a, _, A2 = batched_logdensityof_with_rest(μ.α, A)
+    ℓ_b, _, A_rest = batched_logdensityof_with_rest(μ.β, A2)
+    n_μ = size(A, 1) - size(A_rest, 1)
+    A_μ = view(A, 1:n_μ, Base.tail(axes(A))...)
+    return ℓ_a .+ ℓ_b, A_μ, A_rest
+end
+
 function logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector)
     ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x)
     ℓ_b, b, x_rest = logdensityof_with_rest(μ.β, x2)
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 3a2f77aa..400a6990 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -105,6 +105,7 @@ end
 
 @inline logdensityof_impl(μ::PowerMeasure, x) = _powered_ld(logdensityof_impl, μ, x)
 @inline logdensity_def(μ::PowerMeasure, x) = _powered_ld(logdensity_def, μ, x)
+@inline batched_logdensityof_impl(μ::PowerMeasure, A::AbstractArray) = _batched_ld(logdensityof_impl, μ, A)
 
 # Support checks of powers run over the flat variate storage where the base
 # measure has scalar variates, elementwise otherwise:
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 9d44112b..bdaea979 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -199,6 +199,27 @@ end
 
 marginals(μ::ProductMeasure) = μ.marginals
 
+@inline mspace_elsize(μ::ProductMeasure{<:AbstractArray}) = maybestatic_size(marginals(μ))
+
+@inline function mspace_flatsize(μ::ProductMeasure{<:AbstractArray{M}}) where {M}
+    _cat_sizes(mspace_flatsize(M), maybestatic_size(marginals(μ)))
+end
+
+# The marginals align with the leading dimensions of the flat batch, so
+# one broadcast evaluates all marginal densities:
+@inline function batched_logdensityof_impl(μ::ProductMeasure{<:AbstractArray{M,N}}, A::AbstractArray) where {M,N}
+    _product_batched_ld(μ, A, mspace_flatsize(M), Val(N))
+end
+
+@inline function _product_batched_ld(μ::ProductMeasure, A::AbstractArray, ::Tuple{}, ::Val{N}) where {N}
+    ld = Broadcast.instantiate(Broadcast.broadcasted(dynamic ∘ logdensityof_impl, marginals(μ), A))
+    _sum_leading_dims(ld, static(N))
+end
+
+@inline function _product_batched_ld(μ::ProductMeasure, A::AbstractArray, ::Any, ::Val)
+    _batched_ld_generic(logdensityof_impl, μ, A)
+end
+
 # TODO: Better `map` support in MappedArrays
 _map(f, args...) = map(f, args...)
 _map(f, x::MappedArrays.ReadonlyMappedArray) = mappedarray(fchain((x.f, f)), x.data)
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index 6341e654..64cbb0a8 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -26,6 +26,10 @@ _logweight(::AbstractMeasure) = 0
     _logweight_for(d.logweight, x) + logdensityof_impl(basemeasure(d), x)
 end
 
+@inline function batched_logdensityof_impl(d::AbstractWeightedMeasure, A::AbstractArray)
+    _lazy_add(_logweight_for(d.logweight, A), batched_logdensityof_impl(basemeasure(d), A))
+end
+
 function Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractWeightedMeasure) where {T}
     rand(rng, T, basemeasure(μ))
 end
diff --git a/src/density-batched.jl b/src/density-batched.jl
index 03f7a138..26ddb3c3 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -34,9 +34,11 @@ flat storage: the leading dimensions of `A` are the variate dimensions
 batch dimensions. Returns the log-densities as an array over the batch
 dimensions, or a scalar if there are none.
 
-Power measures never reach `batched_logdensityof_impl`, their power
-structure is unwrapped beforehand. Implementations must handle points
-outside the support of `μ` (the result must be `-Inf` there).
+Implementations must handle points outside the support of `μ` (the
+result must be `-Inf` there). Implementations for structural measures
+evaluate their component measures via `batched_logdensityof_impl` as well,
+power measures route back into the batched core (which unwraps their
+power structure and sums over the power dimensions).
 
 The default implementation broadcasts the log-density over `A` for
 measures with scalar variates and maps it over the variate slices of `A`
@@ -245,3 +247,51 @@ end
 @inline function _sum_dims_seq(A::AbstractArray, ::StaticInteger{N}) where {N}
     _sum_dims_seq(sum(A; dims = N), static(N - 1))
 end
+
+
+@inline _lazy_add(c, x::Number) = c + x
+@inline _lazy_add(c, A) = Broadcast.instantiate(Broadcast.broadcasted(+, c, A))
+
+
+"""
+    MeasureBase.batched_logdensityof_with_rest(μ::AbstractMeasure, A::AbstractArray)
+
+Batched form of [`MeasureBase.logdensityof_with_rest`](@ref) for a batch
+`A` of flat vector streams: the first dimension of `A` runs along the
+streams, all further dimensions are batch dimensions.
+
+Returns a tuple `(ℓ, A_μ, A_rest)` of the log-densities over the batch
+dimensions, the flat variate batch consumed from the streams and the
+unconsumed rest of the streams.
+
+Requires the flat variate size of `μ` to be known, see
+[`MeasureBase.mspace_flatsize`](@ref).
+"""
+function batched_logdensityof_with_rest end
+
+function batched_logdensityof_with_rest(μ::AbstractMeasure, A::AbstractArray)
+    A_μ, A_rest = _batched_consume(A, mspace_flatsize(μ))
+    return _materialize(_batched_ld(logdensityof_impl, μ, A_μ)), A_μ, A_rest
+end
+
+# Consume the leading rows of a batch of streams as a batch of flat variates:
+@inline function _batched_consume(A::AbstractArray, sz::SizeLike)
+    n = size2length(sz)
+    n_stream = size(A, 1)
+    if n_stream < n
+        throw(ArgumentError("Variate streams too short during batched density evaluation"))
+    end
+    batch_axes = Base.tail(axes(A))
+    A_flat = view(A, 1:dynamic(n), batch_axes...)
+    A_rest = view(A, (dynamic(n) + 1):n_stream, batch_axes...)
+    return maybestatic_reshape(A_flat, (_size_dims(sz)..., map(length, batch_axes)...)), A_rest
+end
+
+@inline function _batched_consume(A::AbstractArray, ::Tuple{})
+    batch_axes = Base.tail(axes(A))
+    view(A, 1, batch_axes...), view(A, 2:size(A, 1), batch_axes...)
+end
+
+function _batched_consume(::AbstractArray, sz::NoMSpaceElementSize)
+    throw(ArgumentError("Batched stream consumption requires a variate of known flat size"))
+end
diff --git a/src/density-core.jl b/src/density-core.jl
index 181f147c..c7de098b 100644
--- a/src/density-core.jl
+++ b/src/density-core.jl
@@ -99,10 +99,14 @@ resp. `testvalue` and delegates to `logdensityof_impl`.
 function logdensityof_with_rest end
 
 function logdensityof_with_rest(μ::AbstractMeasure, x::AbstractVector)
-    a, x_rest = _consume_from_stream(x, some_mspace_elsize(μ))
+    a, x_rest = _consume_from_stream(x, _stream_consume_size(μ))
     return logdensityof_impl(μ, a), a, x_rest
 end
 
+@inline _stream_consume_size(μ) = _stream_consume_size(μ, mspace_flatsize(μ))
+@inline _stream_consume_size(μ, sz::SizeLike) = sz
+@inline _stream_consume_size(μ, ::NoMSpaceElementSize) = some_mspace_elsize(μ)
+
 function logdensityof_with_rest(μ::AbstractMeasure, x::NamedTuple)
     a, x_rest = _split_after(x, Val(_mspace_names(μ)))
     return logdensityof_impl(μ, a), a, x_rest
diff --git a/test/logdensities.jl b/test/logdensities.jl
index 25e31328..98d1323b 100644
--- a/test/logdensities.jl
+++ b/test/logdensities.jl
@@ -3,7 +3,7 @@
 using Test
 
 using MeasureBase
-using MeasureBase: logdensities, logdensity_def, StdNormal, StdUniform, Dirac, Lebesgue, LebesgueBase, superpose, weightedmeasure
+using MeasureBase: logdensities, logdensity_def, StdNormal, StdUniform, StdExponential, StdLogistic, Dirac, Lebesgue, LebesgueBase, superpose, weightedmeasure, mcombine, productmeasure
 using ArraysOfArrays: VectorOfSimilarVectors, sliced, flatview
 using StaticArrays: SVector, @SVector, @SMatrix
 using Static: static
@@ -133,6 +133,75 @@ stdnormal_ld(x) = -(x^2 + log2π) / 2
         @test @inferred(logdensityof(Lebesgue()^3, randn(3))) == 0
     end
 
+    @testset "structural batched kernels" begin
+        w = weightedmeasure(log(0.3), StdNormal()^3)
+        X = randn(3, 10)
+        @test @inferred(logdensities(w, X)) ≈ [logdensityof(w, x) for x in eachcol(X)]
+        xw = randn(3)
+        fw(x) = logdensityof(w, x)
+        @test @inferred(fw(xw)) ≈ log(0.3) + sum(stdnormal_ld, xw)
+        @test @allocated(fw(xw)) == 0
+
+        ms = [weightedmeasure(log(i), StdNormal()) for i in 1:4]
+        prod4 = productmeasure(ms)
+        @test @inferred(MeasureBase.mspace_flatsize(prod4)) == (4,)
+        @test @inferred(MeasureBase.mspace_elsize(prod4)) == (4,)
+        xp = randn(4)
+        fp(x) = logdensityof(prod4, x)
+        @test @inferred(fp(xp)) ≈ sum(log(i) + stdnormal_ld(xp[i]) for i in 1:4)
+        @test @allocated(fp(xp)) == 0
+        Xp = randn(4, 7)
+        @test @inferred(logdensities(prod4, Xp)) ≈ [logdensityof(prod4, x) for x in eachcol(Xp)]
+        @test logdensities(prod4, sliced(Xp, 1)) ≈ logdensities(prod4, Xp)
+        @test @inferred(logdensityof(prod4^2, randn(4, 2))) isa Float64
+        Xpp = randn(4, 2, 5)
+        @test logdensities(prod4^2, Xpp) ≈ [logdensityof(prod4^2, Xpp[:, :, i]) for i in 1:5]
+
+        ms2 = reshape([weightedmeasure(log(i), StdUniform()) for i in 1:6], 2, 3)
+        prod23 = productmeasure(ms2)
+        @test @inferred(MeasureBase.mspace_flatsize(prod23)) == (2, 3)
+        x23 = rand(2, 3)
+        @test @inferred(logdensityof(prod23, x23)) ≈ sum(log(i) for i in 1:6)
+        @test logdensities(prod23, rand(2, 3, 4)) ≈ fill(sum(log(i) for i in 1:6), 4)
+
+        mvec = productmeasure([StdNormal()^2, StdNormal()^2])
+        @test @inferred(MeasureBase.mspace_flatsize(mvec)) isa MeasureBase.NoMSpaceElementSize
+    end
+
+    @testset "batched with-rest for combined measures" begin
+        m = mcombine(vcat, StdNormal()^2, StdUniform()^3)
+        @test @inferred(MeasureBase.mspace_flatsize(m)) == (5,)
+        x = vcat(randn(2), rand(3))
+        @test @inferred(logdensityof(m, x)) ≈ sum(stdnormal_ld, x[1:2])
+        X = vcat(randn(2, 6), rand(3, 6))
+        @test @inferred(logdensities(m, X)) ≈ [logdensityof(m, x) for x in eachcol(X)]
+        @test logdensities(m, sliced(X, 1)) ≈ logdensities(m, X)
+        ℓ, A_μ, A_rest = MeasureBase.batched_logdensityof_with_rest(StdNormal()^2, X)
+        @test ℓ ≈ vec(sum(stdnormal_ld.(X[1:2, :]), dims = 1))
+        @test size(A_μ) == (2, 6) && size(A_rest) == (3, 6)
+        @test_throws ArgumentError logdensities(m, vcat(X, rand(1, 6)))
+
+        m3 = mcombine(vcat, StdNormal(), mcombine(vcat, StdExponential()^2, StdLogistic()))
+        @test @inferred(MeasureBase.mspace_flatsize(m3)) == (4,)
+        X3 = vcat(randn(1, 5), rand(2, 5), randn(1, 5))
+        @test @inferred(logdensities(m3, X3)) ≈ [logdensityof(m3, x) for x in eachcol(X3)]
+    end
+
+    @testset "GPU array semantics for structural kernels" begin
+        JLArrays.allowscalar(false)
+        ms = JLArray([weightedmeasure(log(i), StdNormal()) for i in 1:4])
+        prodj = productmeasure(ms)
+        Xj = JLArray(randn(4, 7))
+        ldj = logdensities(prodj, Xj)
+        @test ldj isa JLArray
+        @test Array(ldj) ≈ logdensities(productmeasure(Array(ms)), Array(Xj))
+        mc = mcombine(vcat, StdNormal()^2, StdUniform()^3)
+        Xc = JLArray(vcat(randn(2, 6), rand(3, 6)))
+        ldc = logdensities(mc, Xc)
+        @test ldc isa JLArray
+        @test Array(ldc) ≈ logdensities(mc, Array(Xc))
+    end
+
     @testset "GPU array semantics" begin
         JLArrays.allowscalar(false)
 

From bc864d65e2ed47696d26b9b9e611bc72186d6a1b Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 00:43:59 +0200
Subject: [PATCH 086/122] Rebuild transport on standard-measure extension
 points

Transport now goes through one preferred standard measure per measure
type: measure types implement transport_to_std and transport_from_std
for that type, combinators implement the stream forms
transport_to_std_with_rest and transport_from_std_with_rest, and the
core converts between standard measure types with tail-accurate
transports in log form. The origin-based machinery, the intermediate
heuristics and the multivariate-standard gateways are gone.

The Distributions extension transports univariate continuous
distributions via StdLogistic using log-cdf, log-ccdf, quantile and
complementary quantile, location-scale families via their affine map,
MvNormal via Cholesky factors and Dirichlet via stick breaking.

Created by generative AI.
---
 ext/MeasureBaseChainRulesCoreExt.jl           |   5 +-
 .../MeasureBaseDistributionsExt.jl            |   7 +-
 ext/MeasureBaseDistributionsExt/dirichlet.jl  |   6 +-
 .../dist_vartransform.jl                      |  18 +
 .../distribution_measure.jl                   |   7 +-
 ext/MeasureBaseDistributionsExt/reshaped.jl   |   6 -
 .../standard_dist.jl                          |  12 +-
 ext/MeasureBaseDistributionsExt/standardmv.jl |  14 +-
 ext/MeasureBaseDistributionsExt/univariate.jl | 177 ++++---
 ext/MeasureBaseDistributionsForwardDiffExt.jl |  54 ++-
 ...aseDistributionsForwardDiffPullbacksExt.jl |  15 +-
 ext/MeasureBaseMooncakeExt.jl                 |   3 +-
 src/MeasureBase.jl                            |   4 +-
 src/collection_utils.jl                       |  20 +-
 src/combinators/bind.jl                       |  43 +-
 src/combinators/combined.jl                   |  41 +-
 src/combinators/half.jl                       |   6 +-
 src/combinators/power.jl                      |  98 ++++
 src/combinators/product.jl                    |  81 ++++
 src/combinators/product_transport.jl          | 447 -----------------
 src/combinators/transformedmeasure.jl         |  14 +-
 src/combinators/weighted.jl                   |  16 +-
 src/interface.jl                              |   5 +-
 src/primitives/dirac.jl                       |   5 +
 src/proxies.jl                                |   4 -
 src/standard/stdconvert.jl                    |  57 +++
 src/standard/stdmeasure.jl                    |  85 +++-
 src/standard/stdnormal.jl                     |   4 +-
 src/standard/stdtraits.jl                     |  73 +--
 src/transport.jl                              | 456 +++++++++---------
 src/utils.jl                                  |  10 +-
 test/combinators/bind.jl                      |   6 +-
 test/distributions/test_mooncake.jl           |   4 +-
 test/distributions/test_shape_contract.jl     |   8 +-
 test/distributions/test_transport.jl          |  60 ++-
 test/shape_contract.jl                        |   2 +-
 test/test_mooncake.jl                         |   3 +-
 test/transport.jl                             |  91 +++-
 38 files changed, 965 insertions(+), 1002 deletions(-)
 delete mode 100644 src/combinators/product_transport.jl
 create mode 100644 src/standard/stdconvert.jl

diff --git a/ext/MeasureBaseChainRulesCoreExt.jl b/ext/MeasureBaseChainRulesCoreExt.jl
index 1f457026..cd042496 100644
--- a/ext/MeasureBaseChainRulesCoreExt.jl
+++ b/ext/MeasureBaseChainRulesCoreExt.jl
@@ -73,7 +73,7 @@ end
 
 # = insupport & friends ======================================================
 
-using MeasureBase: check_dof, require_insupport, checked_arg, _checksupport, _origin_depth
+using MeasureBase: check_dof, require_insupport, checked_arg, _checksupport
 
 @inline function ChainRulesCore.rrule(::typeof(_checksupport), cond, result)
     y = _checksupport(cond, result)
@@ -88,9 +88,6 @@ 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
 
diff --git a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
index a230386a..79df0d97 100644
--- a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
+++ b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
@@ -17,12 +17,12 @@ using MeasureBase: StdMeasure, StdUniform, StdExponential, StdLogistic, StdNorma
 using MeasureBase: PowerMeasure, WeightedMeasure, SuperpositionMeasure, PushforwardMeasure
 using MeasureBase: basemeasure, rootmeasure, testvalue, productmeasure, pushfwd, superpose
 using MeasureBase: getdof, checked_arg, massof
-using MeasureBase: transport_to, transport_def, transport_origin, from_origin, to_origin
-using MeasureBase: NoTransportOrigin, NoTransport
+using MeasureBase: transport_to, transport_def, transport_to_std, transport_from_std
 using MeasureBase: Reshape
 using MeasureBase: convert_realtype, _fwddiff, @_adignore
 import MeasureBase:
-    _dist_params_numtype, _trafo_cdf_impl, _trafo_quantile_impl, _trafo_quantile_impl_generic
+    _dist_params_numtype, _trafo_logcdf_impl, _trafo_logccdf_impl,
+    _trafo_quantile_impl, _trafo_cquantile_impl, _dist_quantile, _dist_cquantile
 using MeasureBase: _pushfront, _pushback, _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log
 
 import Distributions
@@ -38,6 +38,7 @@ import StatsFuns
 import PDMats
 
 using IrrationalConstants: log2π, invsqrt2π
+using LogExpFunctions: logistic
 
 using HeterogeneousComputing: real_numtype
 
diff --git a/ext/MeasureBaseDistributionsExt/dirichlet.jl b/ext/MeasureBaseDistributionsExt/dirichlet.jl
index 98693202..90942710 100644
--- a/ext/MeasureBaseDistributionsExt/dirichlet.jl
+++ b/ext/MeasureBaseDistributionsExt/dirichlet.jl
@@ -5,8 +5,6 @@ const DirichletMeasure = AsMeasure{<:Dirichlet}
 MeasureBase.getdof(d::Dirichlet) = length(d) - 1
 MeasureBase.getdof(m::DirichletMeasure) = getdof(m.obj)
 
-MeasureBase.transport_origin(d::Dirichlet) = StdUniform()^getdof(d)
-
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Dirichlet}) = StdUniform
 
 
@@ -18,7 +16,7 @@ end
 
 _a_times_one_minus_b(a::Real, b::Real) = a * (1 - b)
 
-function MeasureBase.from_origin(ν::Dirichlet, x)
+function MeasureBase.transport_from_std(::Type{StdUniform}, ν::Dirichlet, x)
     # See M. J. Betancourt, "Cruising The Simplex: Hamiltonian Monte Carlo and the Dirichlet Distribution",
     # https://arxiv.org/abs/1010.3436
 
@@ -52,7 +50,7 @@ function _dirichlet_variate_to_beta_v(y::AbstractVector{<:Real})
     return beta_v
 end
 
-function MeasureBase.to_origin(ν::Dirichlet, y)
+function MeasureBase.transport_to_std(::Type{StdUniform}, ν::Dirichlet, y)
     @_adignore @argcheck length(ν) == length(y)
     αs = _dropfront(_rev_cumsum(ν.alpha))
     βs = _dropback(ν.alpha)
diff --git a/ext/MeasureBaseDistributionsExt/dist_vartransform.jl b/ext/MeasureBaseDistributionsExt/dist_vartransform.jl
index ceedabe9..b3e45d2b 100644
--- a/ext/MeasureBaseDistributionsExt/dist_vartransform.jl
+++ b/ext/MeasureBaseDistributionsExt/dist_vartransform.jl
@@ -14,3 +14,21 @@ _std_dist_for(::Type{D}, μ::Any) where {D<:_AnyStdDistribution} = _std_dist(_st
 
 MeasureBase.transport_to(::Type{NU}, μ) where {NU<:_AnyStdDistribution} = transport_to(_std_dist_for(NU, μ), μ)
 MeasureBase.transport_to(ν, ::Type{MU}) where {MU<:_AnyStdDistribution} = transport_to(ν, _std_dist_for(MU, ν))
+
+# Disambiguation between the type forms of standard measures and distributions:
+function MeasureBase.transport_to(::Type{NU}, ::Type{MU}) where {NU<:_AnyStdDistribution,MU<:_AnyStdDistribution}
+    _throw_two_std_types()
+end
+function MeasureBase.transport_to(::Type{NU}, ::Type{MU}) where {NU<:StdMeasure,MU<:_AnyStdDistribution}
+    _throw_two_std_types()
+end
+function MeasureBase.transport_to(::Type{NU}, ::Type{MU}) where {NU<:_AnyStdDistribution,MU<:StdMeasure}
+    _throw_two_std_types()
+end
+function _throw_two_std_types()
+    throw(
+        ArgumentError(
+            "Can't construct a transport function between the types of two standard measures, need a measure instance on one side",
+        ),
+    )
+end
diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index 8a99df31..d7a6eb6e 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -75,9 +75,10 @@ end
 @inline MeasureBase.getdof(m::DistributionMeasure{<:ArrayLikeVariate{0}}) = 1
 
 # Delegate transport to the wrapped distribution:
-@inline MeasureBase.transport_origin(m::DistributionMeasure) = m.obj
-@inline MeasureBase.to_origin(::DistributionMeasure, y) = y
-@inline MeasureBase.from_origin(::DistributionMeasure, x) = x
+@inline MeasureBase.transport_to_std(::Type{S}, m::DistributionMeasure, x) where {S<:StdMeasure} =
+    MeasureBase.transport_to_std(S, m.obj, x)
+@inline MeasureBase.transport_from_std(::Type{S}, m::DistributionMeasure, z) where {S<:StdMeasure} =
+    MeasureBase.transport_from_std(S, m.obj, z)
 
 @inline MeasureBase.paramnames(m::DistributionMeasure) = propertynames(m.obj)
 @inline MeasureBase.params(m::DistributionMeasure) = NamedTuple{propertynames(m.obj)}(Distributions.params(m.obj))
diff --git a/ext/MeasureBaseDistributionsExt/reshaped.jl b/ext/MeasureBaseDistributionsExt/reshaped.jl
index 6efde609..8f55c34c 100644
--- a/ext/MeasureBaseDistributionsExt/reshaped.jl
+++ b/ext/MeasureBaseDistributionsExt/reshaped.jl
@@ -2,12 +2,6 @@
 
 MeasureBase.getdof(μ::ReshapedDistribution) = MeasureBase.getdof(μ.dist)
 
-MeasureBase.transport_origin(μ::ReshapedDistribution) = μ.dist
-
-MeasureBase.to_origin(ν::ReshapedDistribution, y) = reshape(y, size(ν.dist))
-
-MeasureBase.from_origin(ν::ReshapedDistribution, x) = reshape(x, ν.dims)
-
 
 function MeasureBase.AbstractMeasure(d::Distributions.ReshapedDistribution)
     orig_dist = d.dist
diff --git a/ext/MeasureBaseDistributionsExt/standard_dist.jl b/ext/MeasureBaseDistributionsExt/standard_dist.jl
index e175655c..8e8963ad 100644
--- a/ext/MeasureBaseDistributionsExt/standard_dist.jl
+++ b/ext/MeasureBaseDistributionsExt/standard_dist.jl
@@ -33,8 +33,6 @@ function Base.show(io::IO, d::StandardDist{D}) where {D}
 end
 
 
-@inline MeasureBase.transport_def(::MU, μ::MU, x) where {MU<:StandardDist{<:Any,0}} = x
-
 for (A, B) in [
     (Uniform, StdUniform),
     (Exponential, StdExponential),
@@ -43,8 +41,10 @@ for (A, B) in [
 ]
     @eval begin
         @inline MeasureBase.preferred_stdmeasure(::Type{<:StandardDist{$A}}) = $B
-        @inline MeasureBase.transport_origin(d::StandardDist{$A,0}) = $B()
-        @inline MeasureBase.transport_origin(d::StandardDist{$A,N}) where {N} = $B()^size(d)
+        @inline MeasureBase.transport_to_std(::Type{$B}, ::StandardDist{$A,0}, x) = x
+        @inline MeasureBase.transport_from_std(::Type{$B}, ::StandardDist{$A,0}, z) = z
+        @inline MeasureBase.transport_to_std(::Type{$B}, ::StandardDist{$A}, x::AbstractArray) = vec(x)
+        @inline MeasureBase.transport_from_std(::Type{$B}, d::StandardDist{$A}, z::AbstractVector) = reshape(z, size(d))
 
         # StandardDist{$A} and $B are equivalent as measures, so convert
         # instead of wrapping:
@@ -61,10 +61,6 @@ for (A, B) in [
     end
 end
 
-@inline MeasureBase.to_origin(ν::StandardDist, y) = y
-@inline MeasureBase.from_origin(ν::StandardDist, x) = x
-
-
 @inline nonstddist(::StandardDist{D,0}) where {D} = D(Distributions.params(D())...)
 @inline function nonstddist(d::StandardDist{D,N}) where {D,N}
     nonstd0 = nonstddist(StandardDist{D}())
diff --git a/ext/MeasureBaseDistributionsExt/standardmv.jl b/ext/MeasureBaseDistributionsExt/standardmv.jl
index 48b9d53e..b21a7ce5 100644
--- a/ext/MeasureBaseDistributionsExt/standardmv.jl
+++ b/ext/MeasureBaseDistributionsExt/standardmv.jl
@@ -4,8 +4,6 @@
 MeasureBase.getdof(d::AbstractMvNormal) = length(d)
 MeasureBase.getdof(m::AsMeasure{<:AbstractMvNormal}) = getdof(m.obj)
 
-MeasureBase.transport_origin(ν::MvNormal) = StandardDist{Normal}(length(ν))
-
 @inline MeasureBase.preferred_stdmeasure(::Type{<:AbstractMvNormal}) = StdNormal
 
 _cholesky_L(A) = cholesky(A).L
@@ -13,16 +11,12 @@ _cholesky_L(A::Diagonal{<:Real}) = Diagonal(sqrt.(diag(A)))
 _cholesky_L(A::PDMats.PDiagMat{<:Real}) = Diagonal(sqrt.(A.diag))
 _cholesky_L(A::PDMats.ScalMat{<:Real}) = Diagonal(Fill(sqrt(A.value), A.dim))
 
-function MeasureBase.from_origin(ν::MvNormal, x)
-    A = _cholesky_L(ν.Σ)
-    b = ν.μ
-    muladd(A, x, b)
+function MeasureBase.transport_to_std(::Type{StdNormal}, d::MvNormal, x)
+    _cholesky_L(d.Σ) \ (x - d.μ)
 end
 
-function MeasureBase.to_origin(ν::MvNormal, y)
-    A = _cholesky_L(ν.Σ)
-    b = ν.μ
-    A \ (y - b)
+function MeasureBase.transport_from_std(::Type{StdNormal}, d::MvNormal, z)
+    muladd(_cholesky_L(d.Σ), z, d.μ)
 end
 
 
diff --git a/ext/MeasureBaseDistributionsExt/univariate.jl b/ext/MeasureBaseDistributionsExt/univariate.jl
index e34ffaa3..8d4b26da 100644
--- a/ext/MeasureBaseDistributionsExt/univariate.jl
+++ b/ext/MeasureBaseDistributionsExt/univariate.jl
@@ -3,48 +3,60 @@
 
 @inline MeasureBase.getdof(::Distribution{Univariate}) = static(1)
 
-@inline MeasureBase.preferred_stdmeasure(::Type{<:Distribution{Univariate,Continuous}}) = StdUniform
-@inline MeasureBase.preferred_stdmeasure(::Type{<:Uniform}) = StdUniform
-@inline MeasureBase.preferred_stdmeasure(::Type{<:Exponential}) = StdExponential
-@inline MeasureBase.preferred_stdmeasure(::Type{<:Logistic}) = StdLogistic
-@inline MeasureBase.preferred_stdmeasure(::Type{<:Normal}) = StdNormal
-@inline MeasureBase.preferred_stdmeasure(::Type{<:Distributions.AffineDistribution{<:Any,<:Any,D}}) where {D} = MeasureBase.preferred_stdmeasure(D)
-
 @inline MeasureBase.check_dof(a::Distribution{Univariate}, b::Distribution{Univariate}) = nothing
 
-
-# Generic transformations to/from StdUniform via cdf/quantile:
-
-
 _dist_params_numtype(d::Distribution) = real_numtype(typeof(Distributions.params(d)))
 
+@inline function _result_numtype(d::Distribution{Univariate}, x::T) where {T<:Number}
+    float(promote_type(T, _dist_params_numtype(d)))
+end
 
-@inline _trafo_cdf(d::Distribution{Univariate,Continuous}, x::Number) =
-    _trafo_cdf_impl(_dist_params_numtype(d), d, x)
-
-@inline _trafo_cdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Number) =
-    Distributions.cdf(d, x)
-
-
-@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, u::Number) =
-    _trafo_quantile_impl(_dist_params_numtype(d), d, u)
-
-@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, u::Number) =
-    _trafo_quantile_impl_generic(d, u)
-
-
-@inline _trafo_quantile_impl_generic(d::Distribution{Univariate,Continuous}, u::Number) =
-    Distributions.quantile(d, u)
 
-# Workaround for Beta dist, current quantile implementation only supports Float64:
-@inline function _trafo_quantile_impl_generic(d::Beta{T}, u::Union{Integer,AbstractFloat}) where {T<:Union{Integer,AbstractFloat}}
-    Distributions.quantile(d, convert(promote_type(Float64, typeof(u)), u))
+# Generic transports between univariate continuous distributions and
+# StdLogistic: the log-cdf and log-ccdf keep both tails accurate on the way
+# to the standard measure, quantile and complementary quantile on the way
+# back. The implementation hooks are specialized for dual numbers in the
+# ForwardDiff extension.
+
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Distribution{Univariate,Continuous}}) = StdLogistic
+
+@inline _trafo_logcdf(d::Distribution{Univariate,Continuous}, x::Number) =
+    _trafo_logcdf_impl(_dist_params_numtype(d), d, x)
+@inline _trafo_logccdf(d::Distribution{Univariate,Continuous}, x::Number) =
+    _trafo_logccdf_impl(_dist_params_numtype(d), d, x)
+@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, p::Number) =
+    _trafo_quantile_impl(_dist_params_numtype(d), d, p)
+@inline _trafo_cquantile(d::Distribution{Univariate,Continuous}, p::Number) =
+    _trafo_cquantile_impl(_dist_params_numtype(d), d, p)
+
+@inline _trafo_logcdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Number) =
+    Distributions.logcdf(d, x)
+@inline _trafo_logccdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Number) =
+    Distributions.logccdf(d, x)
+@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, p::Number) =
+    _dist_quantile(d, p)
+@inline _trafo_cquantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, p::Number) =
+    _dist_cquantile(d, p)
+
+@inline _dist_quantile(d::Distribution{Univariate,Continuous}, p::Number) = Distributions.quantile(d, p)
+@inline _dist_cquantile(d::Distribution{Univariate,Continuous}, p::Number) = Distributions.cquantile(d, p)
+
+# The quantile implementation of Beta only supports Float64:
+const _Float64Compatible = Union{Integer,AbstractFloat}
+@inline function _dist_quantile(d::Beta{<:_Float64Compatible}, p::_Float64Compatible)
+    Distributions.quantile(d, convert(promote_type(Float64, typeof(p)), p))
 end
+@inline function _dist_cquantile(d::Beta{<:_Float64Compatible}, p::_Float64Compatible)
+    Distributions.cquantile(d, convert(promote_type(Float64, typeof(p)), p))
+end
+
+# Rounding errors can push quantiles of truncated distributions slightly
+# outside of their support:
+const _Truncated = Distributions.Truncated{<:Distribution{Univariate,Continuous}}
+@inline _dist_quantile(d::_Truncated, p::Real) = _clamp_to_support(d, Distributions.quantile(d, p))
+@inline _dist_cquantile(d::_Truncated, p::Real) = _clamp_to_support(d, Distributions.cquantile(d, p))
 
-# Workaround for rounding errors that can result in quantile values outside of support of Truncated:
-@inline function _trafo_quantile_impl_generic(d::Distributions.Truncated{<:Distribution{Univariate,Continuous}}, u::Real)
-    x = Distributions.quantile(d, u)
-    T = typeof(x)
+function _clamp_to_support(d::_Truncated, x::T) where {T<:Real}
     min_x = T(minimum(d))
     max_x = T(maximum(d))
     if x < min_x && isapprox(x, min_x, atol = 4 * eps(T))
@@ -56,89 +68,68 @@ end
     end
 end
 
-
-@inline function _result_numtype(d::Distribution{Univariate}, x::T) where {T<:Number}
-    float(promote_type(T, _dist_params_numtype(d)))
-end
-
-
-@inline function MeasureBase.transport_def(::StdUniform, μ::Distribution{Univariate,Continuous}, x)
-    R = _result_numtype(μ, x)
-    y = _trafo_cdf(μ, x)
-    ifelse(Distributions.insupport(μ, x), convert(R, y), convert(R, NaN))
+@inline function MeasureBase.transport_to_std(::Type{StdLogistic}, d::Distribution{Univariate,Continuous}, x)
+    R = _result_numtype(d, x)
+    l = _trafo_logcdf(d, x) - _trafo_logccdf(d, x)
+    ifelse(Distributions.insupport(d, x), convert(R, l), convert(R, NaN))
 end
 
-
-@inline function MeasureBase.transport_def(ν::Distribution{Univariate,Continuous}, ::StdUniform, x::T) where {T}
-    R = _result_numtype(ν, x)
-    TF = float(T)
-    # Avoid x ≈ 0 and x ≈ 1 to avoid infinite variate values for target
-    # distributions with infinite support, keep the quantile argument valid
-    # for out-of-range x (the result is masked to NaN then):
-    clamped_x = clamp(convert(TF, x), zero(TF), one(TF))
-    mod_x = ifelse(x ≈ 0, zero(TF) + eps(TF), ifelse(x ≈ 1, one(TF) - eps(TF), clamped_x))
-    y = _trafo_quantile(ν, mod_x)
-    ifelse((zero(x) <= x) & (x <= one(x)), convert(R, y), convert(R, NaN))
+@inline function MeasureBase.transport_from_std(::Type{StdLogistic}, d::Distribution{Univariate,Continuous}, l)
+    R = _result_numtype(d, l)
+    # From the side that keeps the tail:
+    x = l < zero(l) ? _trafo_quantile(d, logistic(l)) : _trafo_cquantile(d, logistic(-l))
+    convert(R, x)
 end
 
 
-# Use standard measures as transformation origin for scaled/translated equivalents:
+# Location-scale families of standard measures transport by their affine map:
 
-function _origin_to_affine(ν::Distribution{Univariate}, y::T) where {T<:Number}
-    trg_offs, trg_scale = Distributions.location(ν), Distributions.scale(ν)
-    x = muladd(y, trg_scale, trg_offs)
-    convert(_result_numtype(ν, y), x)
+@inline function _affine_to_std(d::Distribution{Univariate}, x::Number)
+    z = (x - Distributions.location(d)) / Distributions.scale(d)
+    convert(_result_numtype(d, x), z)
 end
 
-function _affine_to_origin(μ::Distribution{Univariate}, x::T) where {T<:Number}
-    src_offs, src_scale = Distributions.location(μ), Distributions.scale(μ)
-    y = (x - src_offs) / src_scale
-    convert(_result_numtype(μ, x), y)
+@inline function _std_to_affine(d::Distribution{Univariate}, z::Number)
+    x = muladd(z, Distributions.scale(d), Distributions.location(d))
+    convert(_result_numtype(d, z), x)
 end
 
-for (A, B) in [
+for (D, S) in [
     (Uniform, StdUniform),
     (Logistic, StdLogistic),
     (Normal, StdNormal)
 ]
     @eval begin
-        @inline MeasureBase.transport_origin(::$A) = $B()
-        @inline MeasureBase.to_origin(ν::$A, y) = _affine_to_origin(ν, y)
-        @inline MeasureBase.from_origin(ν::$A, x) = _origin_to_affine(ν, x)
+        @inline MeasureBase.preferred_stdmeasure(::Type{<:$D}) = $S
+        @inline MeasureBase.transport_to_std(::Type{$S}, d::$D, x) = _affine_to_std(d, x)
+        @inline MeasureBase.transport_from_std(::Type{$S}, d::$D, z) = _std_to_affine(d, z)
     end
 end
 
-@inline MeasureBase.transport_origin(::Exponential) = StdExponential()
-@inline MeasureBase.to_origin(ν::Exponential, y) = Distributions.scale(ν) \ y
-@inline MeasureBase.from_origin(ν::Exponential, x) = Distributions.scale(ν) * x
-
-
-# Use the underlying distribution as transformation origin for affine
-# transformed distributions:
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Exponential}) = StdExponential
+@inline MeasureBase.transport_to_std(::Type{StdExponential}, d::Exponential, x) =
+    convert(_result_numtype(d, x), Distributions.scale(d) \ x)
+@inline MeasureBase.transport_from_std(::Type{StdExponential}, d::Exponential, z) =
+    convert(_result_numtype(d, z), Distributions.scale(d) * z)
 
-@inline MeasureBase.transport_origin(d::Distributions.AffineDistribution) = d.ρ
-@inline MeasureBase.from_origin(d::Distributions.AffineDistribution, x) = muladd(d.σ, x, d.μ)
-@inline MeasureBase.to_origin(d::Distributions.AffineDistribution, y) = d.σ \ (y - d.μ)
 
+# Affine transformed distributions transport via the underlying distribution:
 
+const _AffineDist = Distributions.AffineDistribution
 
-# Transform between univariate and single-element power measure
+@inline MeasureBase.preferred_stdmeasure(::Type{<:_AffineDist{<:Any,<:Any,D}}) where {D} =
+    MeasureBase.preferred_stdmeasure(D)
 
-function MeasureBase.transport_def(ν::Distribution{Univariate}, μ::PowerMeasure{<:StdMeasure}, x)
-    return transport_def(ν, μ.parent, only(x))
+@inline function MeasureBase.transport_to_std(::Type{S}, d::_AffineDist, x) where {S<:StdMeasure}
+    transport_to_std(S, d.ρ, d.σ \ (x - d.μ))
 end
-
-function MeasureBase.transport_def(ν::PowerMeasure{<:StdMeasure}, μ::Distribution{Univariate}, x)
-    return Fill(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)...)
+@inline function MeasureBase.transport_from_std(::Type{S}, d::_AffineDist, z) where {S<:StdMeasure}
+    muladd(d.σ, transport_from_std(S, d.ρ, z), d.μ)
 end
-
-
-# Transform between univariate and single-element standard multivariate
-
-function MeasureBase.transport_def(ν::Distribution{Univariate}, μ::StandardDist{D,1}, x) where {D}
-    return transport_def(ν, StandardDist{D}(), only(x))
+# Disambiguation with the generic univariate transports:
+@inline function MeasureBase.transport_to_std(::Type{StdLogistic}, d::_AffineDist, x)
+    transport_to_std(StdLogistic, d.ρ, d.σ \ (x - d.μ))
 end
-
-function MeasureBase.transport_def(ν::StandardDist{D,1}, μ::Distribution{Univariate}, x) where {D}
-    return Fill(transport_def(StandardDist{D}(), μ, only(x)), size(ν)...)
+@inline function MeasureBase.transport_from_std(::Type{StdLogistic}, d::_AffineDist, z)
+    muladd(d.σ, transport_from_std(StdLogistic, d.ρ, z), d.μ)
 end
diff --git a/ext/MeasureBaseDistributionsForwardDiffExt.jl b/ext/MeasureBaseDistributionsForwardDiffExt.jl
index 245c8e60..e929066f 100644
--- a/ext/MeasureBaseDistributionsForwardDiffExt.jl
+++ b/ext/MeasureBaseDistributionsForwardDiffExt.jl
@@ -8,29 +8,57 @@ import ForwardDiff
 
 using Distributions: Distribution, Univariate, Continuous, Beta
 
-@inline function MeasureBase._trafo_cdf_impl(
-    ::Type{<:Union{Integer,AbstractFloat}},
+# Dual-number transports for distributions with plain parameters, via the
+# derivatives of cdf and quantile:
+
+const _PlainParams = Type{<:Union{Integer,AbstractFloat}}
+
+@inline function MeasureBase._trafo_logcdf_impl(
+    ::_PlainParams,
+    d::Distribution{Univariate,Continuous},
+    x::ForwardDiff.Dual{TAG},
+) where {TAG}
+    x_v = ForwardDiff.value(x)
+    lp = Distributions.logcdf(d, x_v)
+    dlp_dx = exp(Distributions.logpdf(d, x_v) - lp)
+    ForwardDiff.Dual{TAG}(lp, dlp_dx * ForwardDiff.partials(x))
+end
+
+@inline function MeasureBase._trafo_logccdf_impl(
+    ::_PlainParams,
     d::Distribution{Univariate,Continuous},
     x::ForwardDiff.Dual{TAG},
 ) where {TAG}
     x_v = ForwardDiff.value(x)
-    u = Distributions.cdf(d, x_v)
-    dudx = Distributions.pdf(d, x_v)
-    ForwardDiff.Dual{TAG}(u, dudx * ForwardDiff.partials(x))
+    lp = Distributions.logccdf(d, x_v)
+    dlp_dx = -exp(Distributions.logpdf(d, x_v) - lp)
+    ForwardDiff.Dual{TAG}(lp, dlp_dx * ForwardDiff.partials(x))
 end
 
 @inline function MeasureBase._trafo_quantile_impl(
-    ::Type{<:Union{Integer,AbstractFloat}},
+    ::_PlainParams,
+    d::Distribution{Univariate,Continuous},
+    p::ForwardDiff.Dual{TAG},
+) where {TAG}
+    p_v = ForwardDiff.value(p)
+    x = MeasureBase._dist_quantile(d, p_v)
+    dx_dp = inv(Distributions.pdf(d, x))
+    ForwardDiff.Dual{TAG}(x, dx_dp * ForwardDiff.partials(p))
+end
+
+@inline function MeasureBase._trafo_cquantile_impl(
+    ::_PlainParams,
     d::Distribution{Univariate,Continuous},
-    u::ForwardDiff.Dual{TAG},
+    p::ForwardDiff.Dual{TAG},
 ) where {TAG}
-    x = MeasureBase._trafo_quantile_impl_generic(d, ForwardDiff.value(u))
-    dxdu = inv(Distributions.pdf(d, x))
-    ForwardDiff.Dual{TAG}(x, dxdu * ForwardDiff.partials(u))
+    p_v = ForwardDiff.value(p)
+    x = MeasureBase._dist_cquantile(d, p_v)
+    dx_dp = -inv(Distributions.pdf(d, x))
+    ForwardDiff.Dual{TAG}(x, dx_dp * ForwardDiff.partials(p))
 end
 
-# Workaround for Beta dist, ForwardDiff doesn't work for parameters:
-@inline MeasureBase._trafo_quantile_impl_generic(d::Beta{T}, u::Real) where {T<:ForwardDiff.Dual} =
-    convert(float(typeof(u)), NaN)
+# The quantile of Beta doesn't support dual parameters:
+@inline MeasureBase._dist_quantile(d::Beta{<:ForwardDiff.Dual}, p::Real) = convert(float(typeof(p)), NaN)
+@inline MeasureBase._dist_cquantile(d::Beta{<:ForwardDiff.Dual}, p::Real) = convert(float(typeof(p)), NaN)
 
 end # module MeasureBaseDistributionsForwardDiffExt
diff --git a/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl b/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl
index 667af023..d7ee12b2 100644
--- a/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl
+++ b/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl
@@ -3,7 +3,7 @@
 module MeasureBaseDistributionsForwardDiffPullbacksExt
 
 import MeasureBase
-using MeasureBase: StdMeasure, transport_def
+using MeasureBase: StdMeasure, transport_to_std, transport_from_std
 
 import Distributions
 using Distributions: Distribution, Univariate
@@ -11,15 +11,12 @@ using Distributions: Distribution, Univariate
 import ChainRulesCore
 using ForwardDiffPullbacks: fwddiff
 
-# Use ForwardDiff for univariate transformations:
-@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::Distribution{Univariate}, μ::Distribution{Univariate}, x::Any)
-    ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x)
+# Use ForwardDiff for univariate transports:
+@inline function ChainRulesCore.rrule(::typeof(transport_to_std), ::Type{S}, d::Distribution{Univariate}, x::Any) where {S<:StdMeasure}
+    ChainRulesCore.rrule(fwddiff(transport_to_std), S, d, x)
 end
-@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::StdMeasure, μ::Distribution{Univariate}, x::Any)
-    ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x)
-end
-@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::Distribution{Univariate}, μ::StdMeasure, x::Any)
-    ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x)
+@inline function ChainRulesCore.rrule(::typeof(transport_from_std), ::Type{S}, d::Distribution{Univariate}, z::Any) where {S<:StdMeasure}
+    ChainRulesCore.rrule(fwddiff(transport_from_std), S, d, z)
 end
 
 end # module MeasureBaseDistributionsForwardDiffPullbacksExt
diff --git a/ext/MeasureBaseMooncakeExt.jl b/ext/MeasureBaseMooncakeExt.jl
index a834bbec..899e2036 100644
--- a/ext/MeasureBaseMooncakeExt.jl
+++ b/ext/MeasureBaseMooncakeExt.jl
@@ -7,7 +7,7 @@ import Mooncake
 using Mooncake: @zero_derivative, MinimalCtx
 
 using MeasureBase: isneginf, isposinf, _adignore_call
-using MeasureBase: check_dof, require_insupport, _origin_depth
+using MeasureBase: check_dof, require_insupport
 
 # Unlike Zygote, Mooncake differentiates the collection utilities
 # (`_pushfront`, etc., mutating code in general), `checked_arg` and
@@ -20,7 +20,6 @@ using MeasureBase: check_dof, require_insupport, _origin_depth
 @zero_derivative MinimalCtx Tuple{typeof(_adignore_call),Any}
 
 @zero_derivative MinimalCtx Tuple{typeof(require_insupport),Any,Any}
-@zero_derivative MinimalCtx Tuple{typeof(_origin_depth),Any}
 @zero_derivative MinimalCtx Tuple{typeof(check_dof),Any,Any}
 
 end # module MeasureBaseMooncakeExt
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 5bf5b4be..7c14fa92 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -184,6 +184,7 @@ include("collection_utils.jl")
 include("smf.jl")
 include("mspace.jl")
 include("getdof.jl")
+include("standard/stdmeasure.jl")
 include("transport.jl")
 include("proxies.jl")
 include("parameterized.jl")
@@ -214,13 +215,12 @@ include("combinators/smart-constructors.jl")
 include("combinators/conditional.jl")
 include("combinators/implicitlymapped.jl")
 
-include("standard/stdmeasure.jl")
 include("standard/stduniform.jl")
 include("standard/stdexponential.jl")
 include("standard/stdlogistic.jl")
 include("standard/stdnormal.jl")
+include("standard/stdconvert.jl")
 include("standard/stdtraits.jl")
-include("combinators/product_transport.jl")
 include("combinators/combined.jl")
 include("combinators/bind.jl")
 include("combinators/half.jl")
diff --git a/src/collection_utils.jl b/src/collection_utils.jl
index 5b4d915a..8ab2274a 100644
--- a/src/collection_utils.jl
+++ b/src/collection_utils.jl
@@ -31,21 +31,11 @@ Base.@propagate_inbounds function _get_or_view(A::AbstractVector, from::IntegerL
 end
 
 Base.@propagate_inbounds function _get_or_view(
-    A::AbstractVector,
-    ::StaticInteger{from},
-    ::StaticInteger{until},
-) where {from,until}
-    SVector{until - from + 1}(view(A, from:until))
-end
-
-# ToDo: Specialize for StaticVector instead of SVector?
-Base.@propagate_inbounds function _get_or_view(
-    A::SVector,
-    from::StaticInteger,
-    until::StaticInteger,
-)
-    # ToDo: Improve implementation:
-    SVector(_get_or_view(Tuple(A), from, until))
+    A::StaticVector,
+    from::StaticInteger{F},
+    until::StaticInteger{U},
+) where {F,U}
+    SVector{U - F + 1,eltype(A)}(_get_or_view(Tuple(A), from, until))
 end
 
 Base.@propagate_inbounds function _get_or_view(tpl::Tuple, from::IntegerLike, until::IntegerLike)
diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl
index f1d293c4..7a2bbed8 100644
--- a/src/combinators/bind.jl
+++ b/src/combinators/bind.jl
@@ -347,18 +347,39 @@ function Base.rand(rng::Random.AbstractRNG, μ::Bind)
 end
 
 
-function transport_to_mvstd(ν_inner::StdMeasure, μ::Bind, ab)
-    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab)
-    β_a = _get_β_a(μ, a)
-    y1 = transport_to_mvstd(ν_inner, tpm_α, a)
-    y2 = transport_to_mvstd(ν_inner, β_a, b)
-    return vcat(y1, y2)
+# Transport consumes the variate parts of the primary and secondary
+# measure in a single pass, analogous to density evaluation:
+
+transport_to_std(::Type{S}, μ::Bind, ab) where {S<:StdMeasure} = _bind_to_std(S, μ.f_c, μ, ab)
+
+function _bind_to_std(::Type{S}, f_c, μ::Bind, ab) where {S}
+    tpm_α, a, b = tpmeasure_split_combined(f_c, μ.α, ab)
+    vcat(_as_stdstream(transport_to_std(S, tpm_α, a)), _as_stdstream(transport_to_std(S, _get_β_a(μ, a), b)))
 end
 
+function _bind_to_std(::Type{S}, ::Union{typeof(vcat),typeof(merge)}, μ::Bind, ab) where {S}
+    z, _, x_rest = transport_to_std_with_rest(S, μ, ab)
+    if !isempty(x_rest)
+        throw(ArgumentError("Variate too long during transport of a bind"))
+    end
+    return z
+end
+
+function transport_to_std_with_rest(::Type{S}, μ::_BindBy{typeof(vcat)}, x::AbstractVector) where {S<:StdMeasure}
+    z_a, a, x2 = transport_to_std_with_rest(S, μ.α, x)
+    z_b, _, x_rest = transport_to_std_with_rest(S, _get_β_a(μ, a), x2)
+    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    return vcat(z_a, z_b), x_μ, x_rest
+end
+
+function transport_to_std_with_rest(::Type{S}, μ::_BindBy{typeof(merge)}, x::NamedTuple) where {S<:StdMeasure}
+    z_a, a, x2 = transport_to_std_with_rest(S, μ.α, x)
+    z_b, b, x_rest = transport_to_std_with_rest(S, _get_β_a(μ, a), x2)
+    return vcat(z_a, z_b), merge(a, b), x_rest
+end
 
-function transport_from_mvstd_with_rest(ν::Bind, μ_inner::StdMeasure, x)
-    a, x2 = transport_from_mvstd_with_rest(ν.α, μ_inner, x)
-    β_a = _get_β_a(ν, a)
-    b, x_rest = transport_from_mvstd_with_rest(β_a, μ_inner, x2)
-    return ν.f_c(a, b), x_rest
+function transport_from_std_with_rest(::Type{S}, μ::Bind, z::AbstractVector) where {S<:StdMeasure}
+    a, z2 = transport_from_std_with_rest(S, μ.α, z)
+    b, z_rest = transport_from_std_with_rest(S, _get_β_a(μ, a), z2)
+    return μ.f_c(a, b), z_rest
 end
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index 9ebe980b..886611f5 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -228,16 +228,39 @@ function Base.rand(rng::Random.AbstractRNG, μ::CombinedMeasure)
 end
 
 
-function transport_to_mvstd(ν_inner::StdMeasure, μ::CombinedMeasure, ab)
-    tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab)
-    y1 = transport_to_mvstd(ν_inner, tpm_α, a)
-    y2 = transport_to_mvstd(ν_inner, μ.β, b)
-    return vcat(y1, y2)
+# Transport consumes the variate parts of both component measures in a
+# single pass, analogous to density evaluation:
+
+transport_to_std(::Type{S}, μ::CombinedMeasure, ab) where {S<:StdMeasure} = _combined_to_std(S, μ.f_c, μ, ab)
+
+function _combined_to_std(::Type{S}, f_c, μ::CombinedMeasure, ab) where {S}
+    tpm_α, a, b = tpmeasure_split_combined(f_c, μ.α, ab)
+    vcat(_as_stdstream(transport_to_std(S, tpm_α, a)), _as_stdstream(transport_to_std(S, μ.β, b)))
+end
+
+function _combined_to_std(::Type{S}, ::Union{typeof(vcat),typeof(merge)}, μ::CombinedMeasure, ab) where {S}
+    z, _, x_rest = transport_to_std_with_rest(S, μ, ab)
+    if !isempty(x_rest)
+        throw(ArgumentError("Variate too long during transport of a combined measure"))
+    end
+    return z
+end
+
+function transport_to_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector) where {S<:StdMeasure}
+    z_a, _, x2 = transport_to_std_with_rest(S, μ.α, x)
+    z_b, _, x_rest = transport_to_std_with_rest(S, μ.β, x2)
+    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    return vcat(z_a, z_b), x_μ, x_rest
 end
 
+function transport_to_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(merge)}, x::NamedTuple) where {S<:StdMeasure}
+    z_a, a, x2 = transport_to_std_with_rest(S, μ.α, x)
+    z_b, b, x_rest = transport_to_std_with_rest(S, μ.β, x2)
+    return vcat(z_a, z_b), merge(a, b), x_rest
+end
 
-function transport_from_mvstd_with_rest(ν::CombinedMeasure, μ_inner::StdMeasure, x)
-    a, x2 = transport_from_mvstd_with_rest(ν.α, μ_inner, x)
-    b, x_rest = transport_from_mvstd_with_rest(ν.β, μ_inner, x2)
-    return ν.f_c(a, b), x_rest
+function transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure, z::AbstractVector) where {S<:StdMeasure}
+    a, z2 = transport_from_std_with_rest(S, μ.α, z)
+    b, z_rest = transport_from_std_with_rest(S, μ.β, z2)
+    return μ.f_c(a, b), z_rest
 end
diff --git a/src/combinators/half.jl b/src/combinators/half.jl
index e4c35788..899bf52f 100644
--- a/src/combinators/half.jl
+++ b/src/combinators/half.jl
@@ -7,7 +7,7 @@ end
 @inline mspace_elsize(μ::Half) = mspace_elsize(μ.parent)
 @inline mspace_flatsize(μ::Half) = mspace_flatsize(μ.parent)
 @inline mspace_flatsize(::Type{<:Half{M}}) where {M} = mspace_flatsize(M)
-@inline preferred_stdmeasure(::Type{MU}) where {MU<:Half} = NoStdTransport{MU}
+@inline preferred_stdmeasure(::Type{<:Half}) = StdUniform
 
 function Base.show(io::IO, μ::Half)
     print(io, "Half")
@@ -49,5 +49,5 @@ function invsmf(μ::Half, p)
     invsmf(μ.parent, (p + 1) / 2)
 end
 
-transport_def(μ::Half, ::StdUniform, p) = invsmf(μ, p)
-transport_def(::StdUniform, μ::Half, x) = smf(μ, x)
+@inline transport_to_std(::Type{StdUniform}, μ::Half, x) = smf(μ, x)
+@inline transport_from_std(::Type{StdUniform}, μ::Half, p) = invsmf(μ, p)
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 400a6990..80bc1c2c 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -161,3 +161,101 @@ checked_arg(μ::PowerMeasure, x::Any) = _throw_size_mismatch()
 massof(m::PowerMeasure) = massof(m.parent)^prod(m.axes)
 
 
+# Transport: the standard variate of a power is the flat vector of the
+# standard variates of its base measure, in the order of the flat variate
+# storage.
+
+function transport_to_std(::Type{S}, μ::PowerMeasure, x::AbstractArray) where {S<:StdMeasure}
+    _pwr_to_std(S, μ, x, _flat_storage(x), mspace_flatsize(μ))
+end
+
+# Flat storage of known flat size: transport the variates of the innermost
+# base measure over the flat storage.
+function _pwr_to_std(::Type{S}, μ::PowerMeasure, x::AbstractArray, x_flat::AbstractArray, sz_flat::SizeLike) where {S}
+    ν, _ = _pwr_unwrap(μ)
+    _check_flatsize(x_flat, sz_flat)
+    _pwr_to_std_flat(S, ν, x_flat, mspace_flatsize(ν))
+end
+
+@inline function _pwr_to_std_flat(::Type{S}, ν, x_flat::AbstractArray, ::Tuple{}) where {S}
+    _flat_std_of(broadcast(Base.Fix1(_ToStd{S}(), ν), x_flat))
+end
+
+@inline function _pwr_to_std_flat(::Type{S}, ν, x_flat::AbstractArray, sz::SizeLike) where {S}
+    _flat_std_of(map(Base.Fix1(_ToStd{S}(), ν), sliced(x_flat, Val(length(sz)))))
+end
+
+# Otherwise transport the variates of the base measure one by one:
+function _pwr_to_std(::Type{S}, μ::PowerMeasure, x::AbstractArray, ::Any, ::Any) where {S}
+    _flat_std_of(map(Base.Fix1(_ToStd{S}(), pwr_base(μ)), x))
+end
+
+function transport_from_std(::Type{S}, μ::PowerMeasure, z::AbstractVector) where {S<:StdMeasure}
+    _check_stdlength(z, fast_dof(μ))
+    _pwr_from_std(S, μ, z, mspace_flatsize(μ))
+end
+
+function _pwr_from_std(::Type{S}, μ::PowerMeasure, z::AbstractVector, sz_flat::SizeLike) where {S}
+    ν, _ = _pwr_unwrap(μ)
+    _pwr_variate(μ, _pwr_from_std_flat(S, ν, z, sz_flat, mspace_flatsize(ν)))
+end
+
+# Base measures of unknown variate size are transported one by one:
+function _pwr_from_std(::Type{S}, μ::PowerMeasure, z::AbstractVector, ::NoMSpaceElementSize) where {S}
+    ys, z_rest = _marginals_from_std_with_rest(S, marginals(μ), z)
+    if !isempty(z_rest)
+        throw(ArgumentError("Length of standard variate doesn't match degrees of freedom of power measure"))
+    end
+    return ys
+end
+
+@inline function _check_stdlength(z::AbstractVector, n::IntegerLike)
+    if maybestatic_length(z) != n
+        throw(ArgumentError("Length of standard variate doesn't match degrees of freedom of measure"))
+    end
+    return nothing
+end
+@inline _check_stdlength(::AbstractVector, ::AbstractNoDOF) = nothing
+
+@inline function _pwr_from_std_flat(::Type{S}, ν, z::AbstractVector, sz_flat, ::Tuple{}) where {S}
+    maybestatic_reshape(broadcast(Base.Fix1(_FromStd{S}(), ν), z), sz_flat)
+end
+
+@inline function _pwr_from_std_flat(::Type{S}, ν, z::AbstractVector, sz_flat, sz_ν::SizeLike) where {S}
+    n_variates = size2length(sz_flat) ÷ size2length(sz_ν)
+    maybestatic_reshape(stacked(_pwr_from_std_chunks(S, ν, z, n_variates, fast_dof(ν))), sz_flat)
+end
+
+function _pwr_from_std_chunks(::Type{S}, ν, z::AbstractVector, n_variates, dof_ν::IntegerLike) where {S}
+    chunks = sliced(maybestatic_reshape(z, (dof_ν, n_variates)), Val(1))
+    map(Base.Fix1(_FromStd{S}(), ν), chunks)
+end
+
+function _pwr_from_std_chunks(::Type{S}, ν, z::AbstractVector, n_variates, ::AbstractNoDOF) where {S}
+    ys, z_rest = _marginals_from_std_with_rest(S, FillArrays.Fill(ν, n_variates), z)
+    if !isempty(z_rest)
+        throw(ArgumentError("Length of standard variate doesn't match degrees of freedom of power measure"))
+    end
+    return ys
+end
+
+# Powers of measures without fast degrees of freedom transport their
+# elements sequentially:
+function transport_from_std_with_rest(::Type{S}, μ::PowerMeasure, z::AbstractVector) where {S<:StdMeasure}
+    _pwr_from_std_with_rest(S, μ, z, fast_dof(μ))
+end
+@inline _pwr_from_std_with_rest(::Type{S}, μ, z, n::IntegerLike) where {S} = _from_std_with_rest_bydof(S, μ, z, n)
+function _pwr_from_std_with_rest(::Type{S}, μ, z, ::AbstractNoDOF) where {S}
+    _marginals_from_std_with_rest(S, marginals(μ), z)
+end
+
+# The nested variate layout of a power over its flat storage:
+@inline _pwr_variate(μ::PowerMeasure, A::AbstractArray) = _pwr_nest(pwr_base(μ), _pwr_variate(pwr_base(μ), A))
+@inline _pwr_variate(ν, A::AbstractArray) = _nest_leaf(A, mspace_flatsize(ν))
+@inline _nest_leaf(A::AbstractArray, ::Tuple{}) = A
+@inline _nest_leaf(A::AbstractArray, ::NoMSpaceElementSize) = A
+@inline _nest_leaf(A::AbstractArray{<:Any,N}, sz::SizeLike) where {N} = _nest_leaf(A, Val(length(sz)), Val(N))
+@inline _nest_leaf(A::AbstractArray, ::Val{N}, ::Val{N}) where {N} = A
+@inline _nest_leaf(A::AbstractArray, ::Val{M}, ::Val) where {M} = sliced(A, Val(M))
+@inline _pwr_nest(ν::PowerMeasure, B::AbstractArray) = sliced(B, Val(length(pwr_axes(ν))))
+@inline _pwr_nest(ν, B::AbstractArray) = B
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index bdaea979..f78fddbc 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -287,3 +287,84 @@ function checked_arg(
 ) where {names}
     NamedTuple{names}(map(checked_arg, values(marginals(μ)), values(x)))
 end
+
+
+# Transport marginal by marginal, the standard variates of the marginals
+# are concatenated in order:
+
+function transport_to_std(::Type{S}, μ::ProductMeasure{<:Tuple}, x::Tuple) where {S<:StdMeasure}
+    _flatten_to_rv(map((m, xi) -> _as_stdstream(transport_to_std(S, m, xi)), marginals(μ), x))
+end
+
+function transport_to_std(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, x::NamedTuple{names}) where {S<:StdMeasure,names}
+    transport_to_std(S, productmeasure(values(marginals(μ))), values(x))
+end
+
+function transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray}, x::AbstractArray) where {S<:StdMeasure}
+    _flat_std_of(broadcast(_ToStd{S}(), marginals(μ), x))
+end
+
+function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:Tuple}, z::AbstractVector) where {S<:StdMeasure}
+    _marginals_from_std_with_rest(S, marginals(μ), z)
+end
+
+function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, z::AbstractVector) where {S<:StdMeasure,names}
+    ys, z_rest = _marginals_from_std_with_rest(S, values(marginals(μ)), z)
+    return NamedTuple{names}(ys), z_rest
+end
+
+function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray}, z::AbstractVector) where {S<:StdMeasure}
+    _array_product_from_std_with_rest(S, μ, z, fast_dof(μ))
+end
+@inline function _array_product_from_std_with_rest(::Type{S}, μ, z, n::IntegerLike) where {S}
+    _from_std_with_rest_bydof(S, μ, z, n)
+end
+function _array_product_from_std_with_rest(::Type{S}, μ, z, ::AbstractNoDOF) where {S}
+    _marginals_from_std_with_rest(S, marginals(μ), z)
+end
+
+# Marginals with scalar variates transport in a single broadcast:
+function transport_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, z::AbstractVector) where {S<:StdMeasure,M}
+    _array_product_from_std(S, μ, z, mspace_flatsize(M))
+end
+function _array_product_from_std(::Type{S}, μ, z::AbstractVector, ::Tuple{}) where {S}
+    mar = marginals(μ)
+    broadcast(_FromStd{S}(), mar, maybestatic_reshape(z, maybestatic_size(mar)))
+end
+function _array_product_from_std(::Type{S}, μ, z::AbstractVector, ::Any) where {S}
+    ys, z_rest = _marginals_from_std_with_rest(S, marginals(μ), z)
+    if !isempty(z_rest)
+        throw(ArgumentError("Length of standard variate doesn't match degrees of freedom of product measure"))
+    end
+    return ys
+end
+
+function _marginals_from_std_with_rest(::Type{S}, νs::Tuple{Vararg{Any}}, z::AbstractVector) where {S}
+    y1, z_rest = transport_from_std_with_rest(S, νs[1], z)
+    y2_end, z_final_rest = _marginals_from_std_with_rest(S, Base.tail(νs), z_rest)
+    return (y1, y2_end...), z_final_rest
+end
+
+_marginals_from_std_with_rest(::Type{S}, ::Tuple{}, z::AbstractVector) where {S} = (), z
+
+function _marginals_from_std_with_rest(::Type{S}, νs::AbstractArray{M}, z::AbstractVector) where {S,M}
+    idxs = eachindex(νs)
+    if isconcretetype(M)
+        # The variate type is uniform, so the loop is type stable (the type
+        # of the remaining stream stays invariant under repeated view-taking):
+        y1, z_rest = transport_from_std_with_rest(S, νs[first(idxs)], z)
+        ys = similar(νs, typeof(y1))
+        ys[first(idxs)] = y1
+        for i in Iterators.drop(idxs, 1)
+            ys[i], z_rest = transport_from_std_with_rest(S, νs[i], z_rest)
+        end
+        return ys, z_rest
+    else
+        ys_any = Vector{Any}(undef, length(idxs))
+        z_rest = z
+        for (j, i) in enumerate(idxs)
+            ys_any[j], z_rest = transport_from_std_with_rest(S, νs[i], z_rest)
+        end
+        return [y for y in ys_any], z_rest
+    end
+end
diff --git a/src/combinators/product_transport.jl b/src/combinators/product_transport.jl
deleted file mode 100644
index a534300b..00000000
--- a/src/combinators/product_transport.jl
+++ /dev/null
@@ -1,447 +0,0 @@
-"""
-    transport_to(ν, ::Type{MU}) where {MU<:StdMeasure}
-    transport_to(::Type{NU}, μ) where {NU<:StdMeasure}
-
-As a user convenience, a standard measure type like [`StdUniform`](@ref),
-[`StdExponential`](@ref), [`StdNormal`](@ref) or [`StdLogistic`](@ref)
-may be used directly as the source or target of a measure transport.
-
-Depending on [`MeasureBase.some_dof(μ)`](@ref) (resp. `ν`), an instance of
-the standard measure itself or a power of it will be automatically chosen as
-the transport partner.
-
-Example:
-
-```julia
-transport_to(StdNormal, μ)
-transport_to(ν, StdNormal)
-```
-"""
-function transport_to(ν, ::Type{MU}) where {MU<:StdMeasure}
-    transport_to(ν, _std_tp_partner(MU, ν))
-end
-
-function transport_to(::Type{NU}, μ) where {NU<:StdMeasure}
-    transport_to(_std_tp_partner(NU, μ), μ)
-end
-
-function transport_to(::Type{NU}, ::Type{MU}) where {NU<:StdMeasure,MU<:StdMeasure}
-    throw(
-        ArgumentError(
-            "Can't construct a transport function between the types of two standard measures, need a measure instance on one side",
-        ),
-    )
-end
-
-_std_tp_partner(::Type{M}, μ) where {M<:StdMeasure} = _std_tp_partner_bydof(M, some_dof(μ))
-_std_tp_partner_bydof(::Type{M}, ::StaticInteger{1}) where {M<:StdMeasure} = M()
-_std_tp_partner_bydof(::Type{M}, dof::IntegerLike) where {M<:StdMeasure} = M()^dof
-function _std_tp_partner_bydof(::Type{M}, ::AbstractNoDOF{MU}) where {M<:StdMeasure,MU}
-    throw(
-        ArgumentError(
-            "Can't determine a standard transport partner for measures of type $(nameof(MU))",
-        ),
-    )
-end
-
-
-# For transport, always pull a multi-dimensional PowerMeasure back to a
-# one-dimensional PowerMeasure first:
-
-const _PowerMeasureRank1{M} = PowerMeasure{M,<:NTuple{1,OneToLike}}
-
-function transport_origin(μ::PowerMeasure)
-    pwr_base(μ)^prod(pwr_size(μ))
-end
-
-function to_origin(μ::PowerMeasure, x)
-    maybestatic_reshape(x, (prod(pwr_size(μ)),))
-end
-
-function from_origin(μ::PowerMeasure, x_origin)
-    # Sanity check, should never fail:
-    @assert x_origin isa AbstractVector
-    maybestatic_reshape(x_origin, pwr_size(μ))
-end
-
-
-# A one-dimensional PowerMeasure has an origin if its parent has an origin:
-
-function transport_origin(μ::_PowerMeasureRank1)
-    _pwr_origin(typeof(μ), transport_origin(pwr_base(μ)), pwr_axes(μ))
-end
-_pwr_origin(::Type{MU}, parent_origin, axes) where {MU} = parent_origin^axes
-_pwr_origin(::Type{MU}, ::NoTransportOrigin, axes) where {MU} = NoTransportOrigin{MU}()
-
-function to_origin(μ::_PowerMeasureRank1, x)
-    to_origin.(Ref(pwr_base(μ)), x)
-end
-
-function from_origin(μ::_PowerMeasureRank1, x_origin)
-    # Sanity check, should never fail:
-    @assert x_origin isa AbstractVector
-    from_origin.(Ref(pwr_base(μ)), x_origin)
-end
-
-
-# Transport between powers of standard measures, of any rank:
-
-function _stdpow_transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, x) where {NU,MU}
-    y = transport_to(pwr_base(ν), pwr_base(μ)).(x)
-    maybestatic_reshape(y, pwr_size(ν))
-end
-
-function _stdpow_transport_def(ν::StdPowerMeasure{MU}, μ::StdPowerMeasure{MU}, x) where {MU}
-    maybestatic_reshape(x, pwr_size(ν))
-end
-
-transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, x) where {NU,MU} =
-    _stdpow_transport_def(ν, μ, x)
-
-# Disambiguation with the mvstd gateway methods below:
-transport_def(ν::StdPowerMeasure{NU,1}, μ::StdPowerMeasure{MU}, x) where {NU,MU} =
-    _stdpow_transport_def(ν, μ, x)
-transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU,1}, x) where {NU,MU} =
-    _stdpow_transport_def(ν, μ, x)
-transport_def(ν::StdPowerMeasure{NU,1}, μ::StdPowerMeasure{MU,1}, x) where {NU,MU} =
-    _stdpow_transport_def(ν, μ, x)
-
-
-# Transport between univariate standard measures and one-dimensional power
-# measures of size one:
-
-function transport_def(ν::StdMeasure, μ::StdPowerMeasure{MU,1}, x) where {MU}
-    return transport_def(ν, pwr_base(μ), only(x))
-end
-
-function transport_def(ν::StdPowerMeasure{NU,1}, μ::StdMeasure, x) where {NU}
-    sz_ν = pwr_size(ν)
-    @assert prod(sz_ν) == 1
-    return maybestatic_fill(transport_def(pwr_base(ν), μ, x), sz_ν)
-end
-
-
-# Transport to a multivariate standard measure from any measure:
-
-function transport_def(ν::StdPowerMeasure{NU,1}, μ::AbstractMeasure, x) where {NU}
-    y = transport_to_mvstd(pwr_base(ν), μ, x)
-    if maybestatic_length(y) != maybestatic_length(ν)
-        throw(ArgumentError("Length of transport target doesn't match variate DOF during transport"))
-    end
-    return y
-end
-
-function transport_to_mvstd(ν_inner::StdMeasure, μ::AbstractMeasure, x)
-    return _to_mvstd_withdof(ν_inner, μ, fast_dof(μ), x)
-end
-
-# For standard measures and their powers specialized `transport_def` methods
-# exist, for other measures the origin-based machinery must be used directly
-# instead of `transport_def`, to prevent infinite dispatch recursion via the
-# gateway methods above:
-const _StdOrStdPowerMeasure = Union{StdMeasure,StdPowerMeasure}
-
-_transport_def_nongateway(ν::_StdOrStdPowerMeasure, μ::_StdOrStdPowerMeasure, x) =
-    transport_def(ν, μ, x)
-function _transport_def_nongateway(ν, μ, x)
-    _transport_between_origins(ν, _origin_depth(ν), _origin_depth(μ), μ, x)
-end
-
-function _to_mvstd_withdof(ν_inner::StdMeasure, μ::AbstractMeasure, dof_μ::IntegerLike, x)
-    _transport_def_nongateway(ν_inner^dof_μ, μ, x)
-end
-
-function _to_mvstd_withdof(ν_inner::StdMeasure, μ::AbstractMeasure, ::AbstractNoDOF, x)
-    _to_mvstd_withorigin(ν_inner, μ, transport_origin(μ), x)
-end
-
-function _to_mvstd_withorigin(ν_inner::StdMeasure, μ::AbstractMeasure, μ_origin, x)
-    x_origin = to_origin(μ, x)
-    transport_to_mvstd(ν_inner, μ_origin, x_origin)
-end
-
-function _to_mvstd_withorigin(ν_inner::StdMeasure, μ::AbstractMeasure, ::NoTransportOrigin, x)
-    throw(
-        ArgumentError(
-            "Don't know how to transport values of type $(nameof(typeof(x))) from $(nameof(typeof(μ))) to a power of $(nameof(typeof(ν_inner)))",
-        ),
-    )
-end
-
-
-# Transport from a multivariate standard measure to any measure:
-
-function transport_def(ν::AbstractMeasure, μ::StdPowerMeasure{MU,1}, x) where {MU}
-    _transport_from_mvstd(ν, pwr_base(μ), x)
-end
-
-function _transport_from_mvstd(ν::AbstractMeasure, μ_inner::StdMeasure, x)
-    y, x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x)
-    if !isempty(x_rest)
-        throw(ArgumentError("Input value too long during transport"))
-    end
-    return y
-end
-
-function transport_from_mvstd_with_rest(ν::AbstractMeasure, μ_inner::StdMeasure, x)
-    dof_ν = fast_dof(ν)
-    return _from_mvstd_with_rest_withdof(ν, dof_ν, μ_inner, x)
-end
-
-function _from_mvstd_with_rest_withdof(
-    ν::AbstractMeasure,
-    dof_ν::IntegerLike,
-    μ_inner::StdMeasure,
-    x,
-)
-    len_x = maybestatic_length(x)
-
-    # Since we can't check the DOF of measures like Bind, we could "run out
-    # of x" if the original x was too short. `transport_to` below will detect
-    # this, but better throw a more informative exception here:
-    if len_x < dof_ν
-        throw(ArgumentError("Variate too short during transport"))
-    end
-
-    x_inner_dof, x_rest = _split_after(x, dof_ν)
-    y = _transport_def_nongateway(ν, μ_inner^dof_ν, x_inner_dof)
-    return y, x_rest
-end
-
-function _from_mvstd_with_rest_withdof(
-    ν::AbstractMeasure,
-    ::AbstractNoDOF,
-    μ_inner::StdMeasure,
-    x,
-)
-    _from_mvstd_with_rest_withorigin(ν, transport_origin(ν), μ_inner, x)
-end
-
-function _from_mvstd_with_rest_withorigin(
-    ν::AbstractMeasure,
-    ν_origin,
-    μ_inner::StdMeasure,
-    x,
-)
-    x_origin, x_rest = transport_from_mvstd_with_rest(ν_origin, μ_inner, x)
-    from_origin(ν, x_origin), x_rest
-end
-
-function _from_mvstd_with_rest_withorigin(
-    ν::AbstractMeasure,
-    ::NoTransportOrigin,
-    μ_inner::StdMeasure,
-    x,
-)
-    throw(
-        ArgumentError(
-            "Don't know how to transport a value of type $(nameof(typeof(x))) from a power of $(nameof(typeof(μ_inner))) to $(nameof(typeof(ν)))",
-        ),
-    )
-end
-
-
-# Transport between a standard measure and Dirac:
-
-@inline transport_from_mvstd_with_rest(ν::Dirac, ::StdMeasure, x::Any) = ν.x, x
-
-@inline transport_to_mvstd(::StdMeasure, ::Dirac, ::Any) = FillArrays.Zeros{Bool}(0)
-
-
-# Pull back from a product over a Fill to a power measure:
-
-@inline transport_origin(μ::ProductMeasure) = _marginals_tp_origin(marginals(μ))
-@inline to_origin(μ::ProductMeasure, x) = _marginals_to_origin(marginals(μ), x)
-@inline from_origin(μ::ProductMeasure, x_origin) =
-    _marginals_from_origin(marginals(μ), x_origin)
-
-_marginals_tp_origin(::Ms) where {Ms} = NoTransportOrigin{ProductMeasure{Ms}}()
-
-_marginals_tp_origin(marginals_μ::FillArrays.Fill) =
-    _fill_value(marginals_μ)^_fill_axes(marginals_μ)
-_marginals_to_origin(::FillArrays.Fill, x) = x
-_marginals_from_origin(::FillArrays.Fill, x_origin) = x_origin
-
-
-# Pull back from a NamedTuple product measure to a Tuple product measure:
-#
-# Maybe ToDo (breaking): For transport between NamedTuple-marginals we could
-# match names where possible, even if given in different order, and transport
-# between the remaining non-matching names in the order given. This may not
-# be worth the additional complexity, though, since transport is typically
-# used with a (power of a) standard measure on one side.
-
-_marginals_tp_origin(marginals_μ::NamedTuple{names}) where {names} =
-    productmeasure(values(marginals_μ))
-_marginals_to_origin(::NamedTuple{names}, x::NamedTuple{names}) where {names} = values(x)
-_marginals_from_origin(::NamedTuple{names}, x_origin::Tuple) where {names} =
-    NamedTuple{names}(x_origin)
-
-
-# Transport between two instances of ProductMeasure:
-
-transport_def(ν::ProductMeasure, μ::ProductMeasure, x) =
-    _marginal_transport_def(marginals(ν), marginals(μ), x)
-
-function _marginal_transport_def(marginals_ν, marginals_μ, x)
-    @assert size(marginals_ν) == size(marginals_μ) == size(x)  # Sanity check, should not fail
-    transport_def.(marginals_ν, marginals_μ, x)
-end
-
-function _marginal_transport_def(
-    marginals_ν::Tuple{Vararg{AbstractMeasure,N}},
-    marginals_μ::Tuple{Vararg{AbstractMeasure,N}},
-    x,
-) where {N}
-    map(transport_def, marginals_ν, marginals_μ, x)
-end
-
-function _marginal_transport_def(
-    marginals_ν::AbstractVector{<:AbstractMeasure},
-    marginals_μ::Tuple{Vararg{AbstractMeasure,N}},
-    x,
-) where {N}
-    _marginal_transport_def(_as_tuple(marginals_ν, Val(N)), marginals_μ, x)
-end
-
-function _marginal_transport_def(
-    marginals_ν::Tuple{Vararg{AbstractMeasure,N}},
-    marginals_μ::AbstractVector{<:AbstractMeasure},
-    x,
-) where {N}
-    _marginal_transport_def(marginals_ν, _as_tuple(marginals_μ, Val(N)), _as_tuple(x, Val(N)))
-end
-
-
-# Transport from a ProductMeasure to a standard measure:
-
-function transport_to_mvstd(ν_inner::StdMeasure, μ::ProductMeasure, x)
-    _marginals_to_mvstd(ν_inner, marginals(μ), x)
-end
-
-struct _TransportToMvStd{NU<:StdMeasure} <: Function end
-(::_TransportToMvStd{NU})(μ, x) where {NU} = transport_to_mvstd(NU(), μ, x)
-
-function _marginals_to_mvstd(::NU, marginals_μ::Tuple, x::Tuple) where {NU<:StdMeasure}
-    _flatten_to_rv(map(_TransportToMvStd{NU}(), marginals_μ, x))
-end
-
-function _marginals_to_mvstd(
-    ν::NU,
-    marginals_μ::NamedTuple{names},
-    x::NamedTuple{names},
-) where {NU<:StdMeasure,names}
-    _marginals_to_mvstd(ν, values(marginals_μ), values(x))
-end
-
-function _marginals_to_mvstd(::NU, marginals_μ, x) where {NU<:StdMeasure}
-    _flatten_to_rv(broadcast(_TransportToMvStd{NU}(), marginals_μ, x))
-end
-
-
-# Transport from a standard measure to a ProductMeasure, with rest:
-
-const _MaybeUnknownDOF = Union{IntegerLike,AbstractNoDOF}
-
-const _KnownDOFs = Union{Tuple{Vararg{IntegerLike,N}} where N,StaticVector{<:IntegerLike}}
-
-function transport_from_mvstd_with_rest(ν::ProductMeasure, μ_inner::StdMeasure, x)
-    νs = marginals(ν)
-    dofs = map(fast_dof, νs)
-    return _marginals_from_mvstd_with_rest(νs, dofs, μ_inner, x)
-end
-
-function transport_from_mvstd_with_rest(
-    ν::ProductMeasure{<:NamedTuple{names}},
-    μ_inner::StdMeasure,
-    x,
-) where {names}
-    ys, x_rest =
-        transport_from_mvstd_with_rest(productmeasure(values(marginals(ν))), μ_inner, x)
-    return NamedTuple{names}(ys), x_rest
-end
-
-function _dof_access_firstidxs(dofs::Tuple{Vararg{IntegerLike,N}}, first_idx) where {N}
-    cumsum((first_idx, dofs[begin:(end-1)]...))
-end
-
-function _dof_access_firstidxs(dofs::AbstractVector{<:IntegerLike}, first_idx)
-    # ToDo: Improve implementation (reduce memory allocations):
-    cumsum(vcat([eltype(dofs)(first_idx)], dofs[begin:(end-1)]))
-end
-
-function _split_x_by_marginals_with_rest(
-    dofs::Union{Tuple,AbstractVector},
-    x::AbstractVector{<:Number},
-)
-    x_idxs = maybestatic_eachindex(x)
-    first_idxs = _dof_access_firstidxs(dofs, maybestatic_first(x_idxs))
-    xs = map((from, n) -> _get_or_view(x, from, from + n - one(n)), first_idxs, dofs)
-    x_rest = _get_or_view(x, first_idxs[end] + dofs[end], maybestatic_last(x_idxs))
-    return xs, x_rest
-end
-
-function _marginals_from_mvstd_with_rest(
-    νs,
-    dofs::_KnownDOFs,
-    μ_inner::StdMeasure,
-    x::AbstractVector{<:Number},
-)
-    xs, x_rest = _split_x_by_marginals_with_rest(dofs, x)
-    μs = map(n -> μ_inner^n, dofs)
-    ys = map(transport_def, νs, μs, xs)
-    return ys, x_rest
-end
-
-function _marginals_from_mvstd_with_rest(
-    νs,
-    dofs,
-    μ_inner::StdMeasure,
-    x::AbstractVector{<:Number},
-)
-    _marginals_from_mvstd_with_rest_nodof(νs, μ_inner, x)
-end
-
-function _marginals_from_mvstd_with_rest_nodof(
-    νs::Tuple{Vararg{AbstractMeasure}},
-    μ_inner::StdMeasure,
-    x::AbstractVector{<:Number},
-)
-    # ToDo: Check for type stability, may need a generated function:
-    y1, x_rest = transport_from_mvstd_with_rest(νs[1], μ_inner, x)
-    y2_end, x_final_rest = _marginals_from_mvstd_with_rest_nodof(Base.tail(νs), μ_inner, x_rest)
-    return (y1, y2_end...), x_final_rest
-end
-
-_marginals_from_mvstd_with_rest_nodof(::Tuple{}, ::StdMeasure, x::AbstractVector{<:Number}) =
-    (), x
-
-function _marginals_from_mvstd_with_rest_nodof(
-    νs::AbstractVector{M},
-    μ_inner::StdMeasure,
-    x::AbstractVector{<:Number},
-) where {M<:AbstractMeasure}
-    if isconcretetype(M)
-        # Marginals of concrete type produce variates of uniform type, so
-        # the loop below is type stable (the type of the remaining stream
-        # stays invariant under repeated view-taking):
-        idxs = eachindex(νs)
-        y1, x_rest = transport_from_mvstd_with_rest(νs[first(idxs)], μ_inner, x)
-        ys = Vector{typeof(y1)}(undef, length(idxs))
-        ys[begin] = y1
-        j = firstindex(ys) + 1
-        for i in Iterators.drop(idxs, 1)
-            ys[j], x_rest = transport_from_mvstd_with_rest(νs[i], μ_inner, x_rest)
-            j += 1
-        end
-        return ys, x_rest
-    else
-        # Fallback for marginals of mixed type:
-        ys_any = Vector{Any}(undef, length(eachindex(νs)))
-        x_rest = x
-        for (i, ν) in zip(eachindex(ys_any), νs)
-            ys_any[i], x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x_rest)
-        end
-        return [y for y in ys_any], x_rest
-    end
-end
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 706784e6..85b6ad0f 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -201,11 +201,17 @@ _pushfwd_dof(::Type{MU}, ::Type{<:Tuple{Any,Real}}, dof) where {MU} = dof
 # Bypass `checked_arg`, would require potentially costly transformation:
 @inline checked_arg(::PushforwardMeasure, x) = x
 
-@inline transport_origin(ν::PushforwardMeasure) = ν.origin
-@inline from_origin(ν::PushforwardMeasure, x) = ν.f(x)
-@inline to_origin(ν::PushforwardMeasure, y) = ν.finv(y)
+# Pushforwards transport via their origin:
+@inline transport_to_std(::Type{S}, ν::PushforwardMeasure, y) where {S<:StdMeasure} =
+    transport_to_std(S, ν.origin, ν.finv(y))
+@inline transport_from_std(::Type{S}, ν::PushforwardMeasure, z) where {S<:StdMeasure} =
+    ν.f(transport_from_std(S, ν.origin, z))
+@inline function transport_from_std_with_rest(::Type{S}, ν::PushforwardMeasure, z::AbstractVector) where {S<:StdMeasure}
+    x, z_rest = transport_from_std_with_rest(S, ν.origin, z)
+    return ν.f(x), z_rest
+end
 
-massof(m::PushforwardMeasure) = massof(transport_origin(m))
+massof(m::PushforwardMeasure) = massof(m.origin)
 
 function Base.rand(rng::AbstractRNG, ::Type{T}, ν::PushforwardMeasure) where {T}
     return ν.f(rand(rng, T, ν.origin))
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index 64cbb0a8..27f3504e 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -68,8 +68,14 @@ gentype(μ::WeightedMeasure) = gentype(μ.base)
 
 insupport(μ::WeightedMeasure, x) = insupport(μ.base, x)
 
-# TODO: Transports must preserve mass
-transport_origin(ν::WeightedMeasure) = ν.base
-
-to_origin(w::WeightedMeasure, y) = y
-from_origin(w::WeightedMeasure, x) = x
+# Weighted measures transport like their base:
+@inline transport_to_std(::Type{S}, μ::AbstractWeightedMeasure, x) where {S<:StdMeasure} =
+    transport_to_std(S, basemeasure(μ), x)
+@inline transport_from_std(::Type{S}, μ::AbstractWeightedMeasure, z) where {S<:StdMeasure} =
+    transport_from_std(S, basemeasure(μ), z)
+@inline transport_to_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, x::AbstractVector) where {S<:StdMeasure} =
+    transport_to_std_with_rest(S, basemeasure(μ), x)
+@inline transport_to_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, x::NamedTuple) where {S<:StdMeasure} =
+    transport_to_std_with_rest(S, basemeasure(μ), x)
+@inline transport_from_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, z::AbstractVector) where {S<:StdMeasure} =
+    transport_from_std_with_rest(S, basemeasure(μ), z)
diff --git a/src/interface.jl b/src/interface.jl
index 31e68ca0..d938d033 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -6,7 +6,7 @@ using Reexport
 
 using MeasureBase: basemeasure_depth, proxy, istrue
 using MeasureBase: insupport, basemeasure_sequence
-using MeasureBase: transport_to, NoTransport
+using MeasureBase: transport_to
 
 using DensityInterface: logdensityof
 using InverseFunctions: inverse
@@ -90,9 +90,8 @@ function test_transport(ν, μ)
 
     @testset "transport_to $μ to $ν" begin
         x = rand(μ)
-        @test !(@inferred(transport_to(ν, μ)(x)) isa NoTransport)
         f = transport_to(ν, μ)
-        y = f(x)
+        y = @inferred f(x)
         @test structisapprox(@inferred(inverse(f)(y)), x)
         @test @inferred(with_logabsdet_jacobian(f, x)) isa Tuple{supertype(y),Real}
         @test @inferred(with_logabsdet_jacobian(inverse(f), y)) isa Tuple{supertype(x),Real}
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 457d2ce1..5b593da4 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -51,3 +51,8 @@ insupport(d::Dirac, x) = x == d.x
     @boundscheck insupport(μ, x) || throw(ArgumentError("Invalid variate for measure"))
     x
 end
+
+# Dirac measures have no degrees of freedom:
+@inline transport_to_std(::Type{S}, ::Dirac, x) where {S<:StdMeasure} = SVector{0,Bool}()
+@inline transport_from_std(::Type{S}, μ::Dirac, z::AbstractVector) where {S<:StdMeasure} = μ.x
+@inline transport_from_std_with_rest(::Type{S}, μ::Dirac, z::AbstractVector) where {S<:StdMeasure} = μ.x, z
diff --git a/src/proxies.jl b/src/proxies.jl
index f2805176..bfdd13c4 100644
--- a/src/proxies.jl
+++ b/src/proxies.jl
@@ -26,10 +26,6 @@ macro useproxy(M)
         @inline $MeasureBase.getdof(μ::$M) = getdof(proxy(μ))
         @inline $MeasureBase.fast_dof(μ::$M) = fast_dof(proxy(μ))
 
-        @inline $MeasureBase.transport_origin(μ::$M) = transport_origin(proxy(μ))
-        @inline $MeasureBase.to_origin(μ::$M, y) = to_origin(proxy(μ), y)
-        @inline $MeasureBase.from_origin(μ::$M, x) = from_origin(proxy(μ), x)
-
         @inline $MeasureBase.localmeasure(μ::$M, x) = localmeasure(proxy(μ), x)
         @inline $MeasureBase.transportmeasure(μ::$M, x) = transportmeasure(proxy(μ), x)
 
diff --git a/src/standard/stdconvert.jl b/src/standard/stdconvert.jl
new file mode 100644
index 00000000..6a5353ac
--- /dev/null
+++ b/src/standard/stdconvert.jl
@@ -0,0 +1,57 @@
+# Direct transports between standard measures, tail-accurate in both
+# directions. Transports via StdUniform lose the upper tail of unbounded
+# measures, since the uniform variate saturates at one.
+
+# Standard normal log-cdf and log-ccdf:
+@inline _normlogcdf(z) = logerfc(-z * invsqrt2) - logtwo
+@inline _normlogccdf(z) = logerfc(z * invsqrt2) - logtwo
+
+# Complementary standard normal cdf, accurate for large positive arguments:
+@inline _normccdf(z) = erfc(z * invsqrt2) / 2
+
+@inline function transport_def(::StdExponential, ::StdNormal, z)
+    ifelse(z < zero(z), -log1p(-Φ(z)), -log(_normccdf(z)))
+end
+
+@inline function transport_def(::StdNormal, ::StdExponential, x)
+    ifelse(x < oftype(x, logtwo), Φinv(-expm1(-x)), -Φinv(exp(-x)))
+end
+
+@inline transport_def(::StdLogistic, ::StdNormal, z) = _normlogcdf(z) - _normlogccdf(z)
+
+@inline function transport_def(::StdNormal, ::StdLogistic, l)
+    ifelse(l < zero(l), Φinv(logistic(l)), -Φinv(logistic(-l)))
+end
+
+@inline transport_def(::StdLogistic, ::StdExponential, x) = log(-expm1(-x)) + x
+
+@inline transport_def(::StdExponential, ::StdLogistic, l) = log1pexp(l)
+
+
+"""
+    MeasureBase.stdconvert(::Type{S}, ::Type{T}, x)
+
+Convert a variate `x` of the standard measure type `T` into a variate of
+the standard measure type `S`, elementwise for arrays.
+"""
+function stdconvert end
+
+@inline stdconvert(::Type{S}, ::Type{S}, x) where {S<:StdMeasure} = x
+@inline stdconvert(::Type{S}, ::Type{T}, x) where {S<:StdMeasure,T<:StdMeasure} = _StdConvert{S,T}()(x)
+
+struct _StdConvert{S,T} <: Function end
+@inline (::_StdConvert{S,T})(x::Number) where {S,T} = transport_def(S(), T(), x)
+@inline (k::_StdConvert)(x::AbstractArray) = broadcast(k, x)
+
+
+"""
+    MeasureBase.StdPowerMeasure{MU<:StdMeasure,N}
+
+The type of an `N`-dimensional power of a standard measure of type `MU`.
+"""
+const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}}
+
+# Powers of standard measures transport directly, by elementwise conversion:
+function transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, x) where {NU<:StdMeasure,MU<:StdMeasure}
+    _pwr_variate(ν, maybestatic_reshape(stdconvert(NU, MU, x), mspace_flatsize(ν)))
+end
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index bb88b0ae..513aad8f 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -4,13 +4,6 @@ StdMeasure(::typeof(rand)) = StdUniform()
 StdMeasure(::typeof(randexp)) = StdExponential()
 StdMeasure(::typeof(randn)) = StdNormal()
 
-"""
-    MeasureBase.StdPowerMeasure{MU<:StdMeasure,N}
-
-The type of an `N`-dimensional power of a standard measure of type `MU`.
-"""
-const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}}
-
 @inline mspace_elsize(::StdMeasure) = ()
 @inline mspace_flatsize(::StdMeasure) = ()
 @inline mspace_flatsize(::Type{<:StdMeasure}) = ()
@@ -18,3 +11,81 @@ const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}}
 @inline check_dof(::StdMeasure, ::StdMeasure) = nothing
 
 @inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x
+
+@inline transport_to_std(::Type{S}, ::S, x) where {S<:StdMeasure} = x
+@inline transport_from_std(::Type{S}, ::S, z) where {S<:StdMeasure} = z
+
+
+"""
+    struct MeasureBase.NoStdTransport{MU}
+
+Indicates that measures of type `MU` can't be transported to or from a
+standard measure.
+"""
+struct NoStdTransport{MU} end
+
+"""
+    struct MeasureBase.AnyStdMeasure
+
+Indicates that any standard measure serves as transport partner, e.g.
+for measures with zero degrees of freedom.
+"""
+struct AnyStdMeasure end
+
+const _StdTransportPartner = Union{Type{<:StdMeasure},Type{AnyStdMeasure},Type{<:NoStdTransport}}
+
+"""
+    MeasureBase.preferred_stdmeasure(μ)::Type
+    MeasureBase.preferred_stdmeasure(::Type{MU})::Type
+
+The type of standard measure that variates of `μ` are transported to and
+from by default.
+
+Returns `MeasureBase.AnyStdMeasure` if any standard measure serves and
+`MeasureBase.NoStdTransport{MU}` if measures of type `MU` have no
+standard-measure transport. Composite measures combine the preferences of
+their components via [`MeasureBase.promote_stdmeasure`](@ref).
+
+Measure types that support transport to and from standard measures should
+specialize the type-based method.
+"""
+function preferred_stdmeasure end
+
+@inline preferred_stdmeasure(μ) = preferred_stdmeasure(typeof(μ))
+@inline preferred_stdmeasure(::Type{MU}) where {MU} = NoStdTransport{MU}
+
+@inline preferred_stdmeasure(::Type{MU}) where {MU<:StdMeasure} = MU
+
+"""
+    MeasureBase.promote_stdmeasure(A::Type, B::Type)::Type
+
+Combine two results of [`MeasureBase.preferred_stdmeasure`](@ref) into
+the preferred standard measure type of a measure composed of both.
+
+Standard measure types promote to the one with the wider range of
+values that remain distinguishable in floating point arithmetic:
+`StdUniform` promotes to any other standard measure type,
+`StdExponential` to `StdLogistic` and `StdNormal`, and `StdLogistic`
+to `StdNormal`.
+"""
+function promote_stdmeasure end
+
+@inline function promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:StdMeasure,B<:StdMeasure}
+    _stdmeasure_rank(A) >= _stdmeasure_rank(B) ? A : B
+end
+
+@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{B}) where {B} = B
+@inline promote_stdmeasure(::Type{A}, ::Type{AnyStdMeasure}) where {A} = A
+@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{AnyStdMeasure}) = AnyStdMeasure
+@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:NoStdTransport,B} = A
+@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A,B<:NoStdTransport} = B
+@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:NoStdTransport,B<:NoStdTransport} = A
+@inline promote_stdmeasure(::Type{A}, ::Type{AnyStdMeasure}) where {A<:NoStdTransport} = A
+@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{B}) where {B<:NoStdTransport} = B
+
+@inline promote_stdmeasure(::Type{A}) where {A} = A
+@inline function promote_stdmeasure(::Type{A}, ::Type{B}, Cs::Vararg{Type,N}) where {A,B,N}
+    promote_stdmeasure(promote_stdmeasure(A, B), Cs...)
+end
+
+@inline _stdmeasure_rank(::Type{<:StdMeasure}) = 0
diff --git a/src/standard/stdnormal.jl b/src/standard/stdnormal.jl
index f636d311..d870d331 100644
--- a/src/standard/stdnormal.jl
+++ b/src/standard/stdnormal.jl
@@ -1,5 +1,5 @@
-using SpecialFunctions: erfc, erfcinv
-using IrrationalConstants: invsqrt2, log2π
+using SpecialFunctions: erfc, erfcinv, logerfc
+using IrrationalConstants: invsqrt2, log2π, logtwo
 
 struct StdNormal <: StdMeasure end
 
diff --git a/src/standard/stdtraits.jl b/src/standard/stdtraits.jl
index 1562c3a5..d9b7a235 100644
--- a/src/standard/stdtraits.jl
+++ b/src/standard/stdtraits.jl
@@ -1,76 +1,5 @@
-"""
-    struct MeasureBase.NoStdTransport{MU}
+# Standard measure preferences of the combinators and standard measure ranks:
 
-Indicates that measures of type `MU` can't be transported to or from a
-standard measure.
-"""
-struct NoStdTransport{MU} end
-
-"""
-    struct MeasureBase.AnyStdMeasure
-
-Indicates that any standard measure serves as transport partner, e.g.
-for measures with zero degrees of freedom.
-"""
-struct AnyStdMeasure end
-
-const _StdTransportPartner = Union{Type{<:StdMeasure},Type{AnyStdMeasure},Type{<:NoStdTransport}}
-
-"""
-    MeasureBase.preferred_stdmeasure(μ)::Type
-    MeasureBase.preferred_stdmeasure(::Type{MU})::Type
-
-The type of standard measure that variates of `μ` are transported to and
-from by default.
-
-Returns `MeasureBase.AnyStdMeasure` if any standard measure serves and
-`MeasureBase.NoStdTransport{MU}` if measures of type `MU` have no
-standard-measure transport. Composite measures combine the preferences of
-their components via [`MeasureBase.promote_stdmeasure`](@ref).
-
-Measure types that support transport to and from standard measures should
-specialize the type-based method.
-"""
-function preferred_stdmeasure end
-
-@inline preferred_stdmeasure(μ) = preferred_stdmeasure(typeof(μ))
-@inline preferred_stdmeasure(::Type{MU}) where {MU} = NoStdTransport{MU}
-
-@inline preferred_stdmeasure(::Type{MU}) where {MU<:StdMeasure} = MU
-
-"""
-    MeasureBase.promote_stdmeasure(A::Type, B::Type)::Type
-
-Combine two results of [`MeasureBase.preferred_stdmeasure`](@ref) into
-the preferred standard measure type of a measure composed of both.
-
-Standard measure types promote to the one with the wider range of
-values that remain distinguishable in floating point arithmetic:
-`StdUniform` promotes to any other standard measure type,
-`StdExponential` to `StdLogistic` and `StdNormal`, and `StdLogistic`
-to `StdNormal`.
-"""
-function promote_stdmeasure end
-
-@inline function promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:StdMeasure,B<:StdMeasure}
-    _stdmeasure_rank(A) >= _stdmeasure_rank(B) ? A : B
-end
-
-@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{B}) where {B} = B
-@inline promote_stdmeasure(::Type{A}, ::Type{AnyStdMeasure}) where {A} = A
-@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{AnyStdMeasure}) = AnyStdMeasure
-@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:NoStdTransport,B} = A
-@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A,B<:NoStdTransport} = B
-@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:NoStdTransport,B<:NoStdTransport} = A
-@inline promote_stdmeasure(::Type{A}, ::Type{AnyStdMeasure}) where {A<:NoStdTransport} = A
-@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{B}) where {B<:NoStdTransport} = B
-
-@inline promote_stdmeasure(::Type{A}) where {A} = A
-@inline function promote_stdmeasure(::Type{A}, ::Type{B}, Cs::Vararg{Type,N}) where {A,B,N}
-    promote_stdmeasure(promote_stdmeasure(A, B), Cs...)
-end
-
-@inline _stdmeasure_rank(::Type{<:StdMeasure}) = 0
 @inline _stdmeasure_rank(::Type{StdUniform}) = 1
 @inline _stdmeasure_rank(::Type{StdExponential}) = 2
 @inline _stdmeasure_rank(::Type{StdLogistic}) = 3
diff --git a/src/transport.jl b/src/transport.jl
index 65c5f386..0c88e0b6 100644
--- a/src/transport.jl
+++ b/src/transport.jl
@@ -1,51 +1,3 @@
-"""
-    struct MeasureBase.NoTransportOrigin{NU}
-
-Indicates that no (default) pullback measure is available for measures of
-type `NU`.
-
-See [`MeasureBase.transport_origin`](@ref).
-"""
-struct NoTransportOrigin{NU} end
-
-Base.:^(origin::NoTransportOrigin, ::IntegerLike) = origin
-
-"""
-    MeasureBase.transport_origin(ν)
-
-Default measure to pullback to resp. pushforward from when transforming
-between `ν` and another measure.
-"""
-function transport_origin end
-
-transport_origin(ν::NU) where {NU} = NoTransportOrigin{NU}()
-
-"""
-    MeasureBase.from_origin(ν, x)
-
-Push `x` from `MeasureBase.transport_origin(μ)` forward to `ν`.
-"""
-function from_origin end
-
-from_origin(ν::NU, ::Any) where {NU} = NoTransportOrigin{NU}()
-
-"""
-    MeasureBase.to_origin(ν, y)
-
-Pull `y` from `ν` back to `MeasureBase.transport_origin(ν)`.
-"""
-function to_origin end
-
-to_origin(ν::NU, ::Any) where {NU} = NoTransportOrigin{NU}()
-
-"""
-    struct MeasureBase.NoTransport{NU,MU} end
-
-Indicates that no transformation from a measure of type `MU` to a measure of
-type `NU` could be found.
-"""
-struct NoTransport{NU,MU} end
-
 """
     f = transport_to(ν, μ)
 
@@ -54,194 +6,49 @@ Generates a [measurable function](https://en.wikipedia.org/wiki/Measurable_funct
 a value `y = f(x)` distributed according to a measure `ν`.
 
 The [pushforward measure](https://en.wikipedia.org/wiki/Pushforward_measure)
-from `μ` under `f` is is equivalent to `ν`.
-
-If terms of random values this implies that `f(rand(μ))` is equivalent to
-`rand(ν)` (if `rand(μ)` and `rand(ν)` are supported).
-
-The resulting function `f` should support
-`ChangesOfVariables.with_logabsdet_jacobian(f, x)` if mathematically well-defined,
-so that densities of `ν` can be derived from densities of `μ` via `f` (using
-appropriate base measures).
-
-Returns NoTransportOrigin{typeof(ν),typeof(μ)} if no transformation from
-`μ` to `ν` can be found.
-
-To add transformation rules for a measure type `MyMeasure`, specialize
-
-* `MeasureBase.transport_def(ν::SomeStdMeasure, μ::CustomMeasure, x) = ...`
-* `MeasureBase.transport_def(ν::MyMeasure, μ::SomeStdMeasure, x) = ...`
+from `μ` under `f` is equivalent to `ν`, so `f(rand(μ))` is equivalent
+to `rand(ν)`. `f` supports `InverseFunctions.inverse` and
+`ChangesOfVariables.with_logabsdet_jacobian`.
 
-and/or
-
-* `MeasureBase.transport_origin(ν::MyMeasure) = SomeMeasure(...)`
-* `MeasureBase.from_origin(μ::MyMeasure, x) = y`
-* `MeasureBase.to_origin(μ::MyMeasure, y) = x`
-
-and ensure `MeasureBase.getdof(μ::MyMeasure)` is defined correctly.
-
-A standard measure type like `StdUniform`, `StdExponential` or
-`StdLogistic` may also be used as the source or target of the transform:
+Measures are transported via standard measures: `x` is transported to the
+standard measure type that the preferences of `ν` and `μ` promote to (see
+[`MeasureBase.preferred_stdmeasure`](@ref)) and from there to `ν`. A
+standard measure type like `StdUniform` or `StdNormal` may also be used
+directly as the source or target:
 
 ```julia
-f_to_uniform(StdUniform, μ)
-f_to_uniform(ν, StdUniform)
+transport_to(StdNormal, μ)
+transport_to(ν, StdNormal)
 ```
 
-Depending on [`getdof(μ)`](@ref) (resp. `ν`), an instance of the standard
-distribution itself or a power of it (e.g. `StdUniform()` or
-`StdUniform()^dof`) will be chosen as the transformation partner.
+An instance of the standard measure itself or a power of it (depending on
+[`getdof(μ)`](@ref) resp. `ν`) is chosen as the transport partner then.
+
+# Extended help
+
+To support transport for a measure type, specialize
+[`MeasureBase.transport_to_std`](@ref) and
+[`MeasureBase.transport_from_std`](@ref) for its preferred standard measure
+type. Measures whose variates are composed of the variates of other
+measures specialize the stream forms
+[`MeasureBase.transport_to_std_with_rest`](@ref) and
+[`MeasureBase.transport_from_std_with_rest`](@ref) instead.
+[`MeasureBase.transport_def`](@ref) may be specialized for pairs of
+measure types with a direct transport.
 """
 function transport_to end
+export transport_to
 
 """
     transport_to(ν, μ, x)
 
-Transport `x` from the measure `μ` to the measure `ν`
+Transport `x` from the measure `μ` to the measure `ν`, equivalent to
+`transport_to(ν, μ)(x)`.
 """
 transport_to(ν, μ, x) = transport_to(ν, μ)(x)
 
 """
-    transport_def(ν, μ, x)
-
-Transforms a value `x` distributed according to `μ` to a value `y` distributed
-according to `ν`.
-
-If no specialized `transport_def(::MU, ::NU, ...)` is available then
-the default implementation of`transport_def(ν, μ, x)` uses the following
-strategy:
-
-* Evaluate [`transport_origin`](@ref) for μ and ν. Transform between
-  each and it's origin, if available, and use the origin(s) as intermediate
-  measures for another transformation.
-
-* If all else fails, try to transform from μ to a standard multivariate
-  uniform measure and then to ν.
-
-See [`transport_to`](@ref).
-"""
-function transport_def end
-
-function transport_def(ν, μ, x)
-    _transport_between_origins(ν, _origin_depth(ν), _origin_depth(μ), μ, x)
-end
-
-@inline function _origin_depth(ν::NU) where {NU}
-    ν_0 = ν
-    Base.Cartesian.@nexprs 10 i -> begin  # 10 is just some "big enough" number
-        ν_{i} = transport_origin(ν_{i - 1})
-        if ν_{i} isa NoTransportOrigin
-            return static(i - 1)
-        end
-    end
-    return static(10)
-end
-
-# If both measures have no origin:
-function _transport_between_origins(ν, ::StaticInteger{0}, ::StaticInteger{0}, μ, x)
-    _transport_between_noorigins(ν, fast_dof(ν), fast_dof(μ), μ, x)
-end
-
-function _transport_between_noorigins(ν, ::IntegerLike, ::IntegerLike, μ, x)
-    _transport_with_intermediate(ν, _transport_intermediate(ν, μ), μ, x)
-end
-
-# If the DOF of either side is not known, pivot through a flat vector of
-# standard-measure variates, using the with-rest transport protocol:
-_transport_between_noorigins(ν, ::AbstractNoDOF, ::IntegerLike, μ, x) =
-    _transport_mvstd_pivot(ν, μ, x)
-_transport_between_noorigins(ν, ::IntegerLike, ::AbstractNoDOF, μ, x) =
-    _transport_mvstd_pivot(ν, μ, x)
-_transport_between_noorigins(ν, ::AbstractNoDOF, ::AbstractNoDOF, μ, x) =
-    _transport_mvstd_pivot(ν, μ, x)
-
-function _transport_mvstd_pivot(ν, μ, x)
-    z = transport_to_mvstd(StdUniform(), μ, x)
-    return _transport_from_mvstd(ν, StdUniform(), z)
-end
-
-@generated function _transport_between_origins(
-    ν,
-    ::StaticInteger{n_ν},
-    ::StaticInteger{n_μ},
-    μ,
-    x,
-) where {n_ν,n_μ}
-    if n_ν == 10
-        return :(throw(
-            ArgumentError(
-                "Transport to measure of type $(nameof(typeof(ν))) not supported, origin stack too deep.",
-            ),
-        ))
-    end
-    if n_μ == 10
-        return :(throw(
-            ArgumentError(
-                "Transport from measure of type $(nameof(typeof(μ))) not supported, origin stack too deep.",
-            ),
-        ))
-    end
-
-    prog = quote
-        μ0 = μ
-        x0 = x
-        ν0 = ν
-    end
-    for i in 1:n_μ
-        μ_i = Symbol(:μ, i)
-        μ_last = Symbol(:μ, i - 1)
-        push!(prog.args, :($μ_i = transport_origin($μ_last)))
-    end
-    for i in 1:n_μ
-        x_i = Symbol(:x, i)
-        x_last = Symbol(:x, i - 1)
-        μ_last = Symbol(:μ, i - 1)
-        push!(prog.args, :($x_i = to_origin($μ_last, $x_last)))
-    end
-    for i in 1:(n_ν)
-        ν_i = Symbol(:ν, i)
-        ν_last = Symbol(:ν, i - 1)
-        push!(prog.args, :($ν_i = transport_origin($ν_last)))
-    end
-    μ_im = Symbol(:μ, n_μ)
-    x_im = Symbol(:x, n_μ)
-    ν_im = Symbol(:ν, n_ν)
-    y_im = Symbol(:y, n_ν)
-    push!(prog.args, :($y_im = transport_def($ν_im, $μ_im, $x_im)))
-    for i in (n_ν-1):-1:0
-        y_i = Symbol(:y, i)
-        y_last = Symbol(:y, i + 1)
-        ν_last = Symbol(:ν, i)
-        push!(prog.args, :($y_i = from_origin($ν_last, $y_last)))
-    end
-    push!(prog.args, :(return y0))
-    return prog
-end
-
-@inline _transport_intermediate(ν, μ) = _transport_intermediate(fast_dof(ν), fast_dof(μ))
-@inline _transport_intermediate(::IntegerLike, n_μ::IntegerLike) = StdUniform()^n_μ
-@inline _transport_intermediate(::StaticInteger{1}, ::StaticInteger{1}) = StdUniform()
-
-_call_transport_def(ν, μ, x) = transport_def(ν, μ, x)
-_call_transport_def(::Any, ::Any, x::NoTransportOrigin) = x
-_call_transport_def(::Any, ::Any, x::NoTransport) = x
-
-function _transport_with_intermediate(ν, m, μ, x)
-    z = _call_transport_def(m, μ, x)
-    y = _call_transport_def(ν, m, z)
-    return y
-end
-
-# Prevent infinite recursion in case vartransform_intermediate doesn't change type:
-@inline function _transport_with_intermediate(::NU, ::NU, ::MU, ::Any) where {NU,MU}
-    NoTransport{NU,MU}()
-end
-@inline function _transport_with_intermediate(::NU, ::MU, ::MU, ::Any) where {NU,MU}
-    NoTransport{NU,MU}()
-end
-
-"""
-    struct TransportFunction <: Function
+    struct MeasureBase.TransportFunction <: Function
 
 Transforms a variate from one measure to a variate of another.
 
@@ -269,7 +76,7 @@ function Base.:(==)(a::TransportFunction, b::TransportFunction)
 end
 
 Base.@propagate_inbounds function (f::TransportFunction)(x)
-    return _call_transport_def(f.ν, f.μ, checked_arg(f.μ, x))
+    return transport_def(f.ν, f.μ, checked_arg(f.μ, x))
 end
 
 @inline function InverseFunctions.inverse(f::TransportFunction{NU,MU}) where {NU,MU}
@@ -278,11 +85,11 @@ end
 
 function ChangesOfVariables.with_logabsdet_jacobian(f::TransportFunction, x)
     y = f(x)
-    logpdf_src = logdensityof(f.μ, x)
-    logpdf_trg = logdensityof(f.ν, y)
-    ladj = logpdf_src - logpdf_trg
-    # If logpdf_src and logpdf_trg are -Inf setting lafj to zero is safe:
-    fixed_ladj = logpdf_src == logpdf_trg == -Inf ? zero(ladj) : ladj
+    logd_src = logdensityof(f.μ, x)
+    logd_trg = logdensityof(f.ν, y)
+    ladj = logd_src - logd_trg
+    # Both densities being -Inf leaves the Jacobian undefined, zero is a safe choice then:
+    fixed_ladj = ifelse(isneginf(logd_src) & isneginf(logd_trg), zero(ladj), ladj)
     return y, fixed_ladj
 end
 
@@ -309,3 +116,198 @@ function Base.show(io::IO, f::TransportFunction)
 end
 
 Base.show(io::IO, M::MIME"text/plain", f::TransportFunction) = show(io, f)
+
+
+"""
+    MeasureBase.transport_def(ν, μ, x)
+
+Transport a variate `x` of `μ` to a variate of `ν`.
+
+The default implementation transports `x` via the standard measure type the
+preferences of `ν` and `μ` promote to. Specialize `transport_def` for pairs
+of measure types with a direct transport.
+"""
+function transport_def end
+
+@inline transport_def(ν, μ, x) = _transport_via_std(_transport_pivot(ν, μ), ν, μ, x)
+
+function _transport_via_std(::Type{S}, ν, μ, x) where {S<:StdMeasure}
+    z = transport_to_std(S, μ, x)
+    y, z_rest = transport_from_std_with_rest(S, ν, _as_stdstream(z))
+    if !isempty(z_rest)
+        throw(ArgumentError("Degrees of freedom of source and target measure of a transport don't match"))
+    end
+    return y
+end
+
+@inline function _transport_pivot(ν, μ)
+    _concrete_pivot(promote_stdmeasure(preferred_stdmeasure(ν), preferred_stdmeasure(μ)), ν, μ)
+end
+@inline _concrete_pivot(::Type{S}, ν, μ) where {S<:StdMeasure} = S
+@inline _concrete_pivot(::Type{AnyStdMeasure}, ν, μ) = StdUniform
+function _concrete_pivot(::Type{<:NoStdTransport{MU}}, ν, μ) where {MU}
+    throw(ArgumentError("No transport between measures of type $(nameof(typeof(ν))) and $(nameof(typeof(μ))), measures of type $(nameof(MU)) have no transport via standard measures"))
+end
+
+# Standard variates of scalar-variate measures are scalars, streams of
+# standard variates are vectors:
+@inline _as_stdstream(z::AbstractVector) = z
+@inline _as_stdstream(z::Number) = SVector(z)
+
+
+"""
+    MeasureBase.transport_to_std(::Type{S}, μ, x)
+
+Transport a variate `x` of `μ` to a variate of the standard measure type
+`S`: a number if the variates of `μ` are scalars, a flat vector of length
+[`getdof(μ)`](@ref) otherwise.
+
+Measure types specialize `transport_to_std` for their preferred standard
+measure type (see [`MeasureBase.preferred_stdmeasure`](@ref)), the generic
+implementation converts between standard measure types.
+"""
+function transport_to_std end
+
+@inline function transport_to_std(::Type{S}, μ, x) where {S<:StdMeasure}
+    _to_std_via(S, preferred_stdmeasure(μ), μ, x)
+end
+
+@inline function _to_std_via(::Type{S}, ::Type{T}, μ, x) where {S<:StdMeasure,T<:StdMeasure}
+    stdconvert(S, T, transport_to_std(T, μ, x))
+end
+function _to_std_via(::Type{S}, ::Type{S}, μ, x) where {S<:StdMeasure}
+    throw(ArgumentError("Transport to $(nameof(S)) is not implemented for measures of type $(nameof(typeof(μ)))"))
+end
+function _to_std_via(::Type{S}, ::Type, μ, x) where {S<:StdMeasure}
+    throw(ArgumentError("Measures of type $(nameof(typeof(μ))) have no transport via standard measures"))
+end
+
+
+"""
+    MeasureBase.transport_from_std(::Type{S}, μ, z)
+
+Transport a variate `z` of the standard measure type `S` to a variate of
+`μ`, the inverse of [`MeasureBase.transport_to_std`](@ref).
+"""
+function transport_from_std end
+
+@inline function transport_from_std(::Type{S}, μ, z) where {S<:StdMeasure}
+    _from_std_via(S, preferred_stdmeasure(μ), μ, z)
+end
+
+@inline function _from_std_via(::Type{S}, ::Type{T}, μ, z) where {S<:StdMeasure,T<:StdMeasure}
+    transport_from_std(T, μ, stdconvert(T, S, z))
+end
+function _from_std_via(::Type{S}, ::Type{S}, μ, z) where {S<:StdMeasure}
+    throw(ArgumentError("Transport from $(nameof(S)) is not implemented for measures of type $(nameof(typeof(μ)))"))
+end
+function _from_std_via(::Type{S}, ::Type, μ, z) where {S<:StdMeasure}
+    throw(ArgumentError("Measures of type $(nameof(typeof(μ))) have no transport via standard measures"))
+end
+
+
+"""
+    MeasureBase.transport_to_std_with_rest(::Type{S}, μ, x)
+
+Transport the variate of `μ` at the beginning of the stream `x` of
+variate content to the standard measure type `S`.
+
+Returns a tuple `(z, x_μ, x_rest)` of the flat vector `z` of standard
+variates, the variate `x_μ` of `μ` consumed from the stream and the
+unconsumed rest of the stream. See
+[`MeasureBase.logdensityof_with_rest`](@ref) for the stream conventions.
+"""
+function transport_to_std_with_rest end
+
+function transport_to_std_with_rest(::Type{S}, μ, x::AbstractVector) where {S<:StdMeasure}
+    x_μ, x_rest = _consume_from_stream(x, _stream_consume_size(μ))
+    return _as_stdstream(transport_to_std(S, μ, x_μ)), x_μ, x_rest
+end
+
+function transport_to_std_with_rest(::Type{S}, μ, x::NamedTuple) where {S<:StdMeasure}
+    x_μ, x_rest = _split_after(x, Val(_mspace_names(μ)))
+    return _as_stdstream(transport_to_std(S, μ, x_μ)), x_μ, x_rest
+end
+
+
+"""
+    MeasureBase.transport_from_std_with_rest(::Type{S}, μ, z)
+
+Transport the beginning of the flat stream `z` of standard variates of type
+`S` to a variate of `μ`, consuming as many entries as `μ` requires.
+
+Returns a tuple `(x, z_rest)` of the variate `x` and the unconsumed rest of
+the stream. Measure types whose degrees of freedom depend on variate values
+implement `transport_from_std_with_rest` instead of
+[`MeasureBase.transport_from_std`](@ref).
+"""
+function transport_from_std_with_rest end
+
+function transport_from_std_with_rest(::Type{S}, μ, z::AbstractVector) where {S<:StdMeasure}
+    _from_std_with_rest_bydof(S, μ, z, fast_dof(μ))
+end
+
+function _from_std_with_rest_bydof(::Type{S}, μ, z::AbstractVector, n::IntegerLike) where {S}
+    if maybestatic_length(z) < n
+        throw(ArgumentError("Stream of standard variates too short during transport"))
+    end
+    z_μ, z_rest = _split_after(z, n)
+    return transport_from_std(S, μ, _chunk_as_variate(μ, z_μ)), z_rest
+end
+
+function _from_std_with_rest_bydof(::Type{S}, μ, z::AbstractVector, ::AbstractNoDOF) where {S}
+    throw(ArgumentError("Transport from standard measures requires measures of type $(nameof(typeof(μ))) to implement MeasureBase.transport_from_std_with_rest"))
+end
+
+# Scalar-variate measures take their standard variate as a number:
+@inline _chunk_as_variate(μ, z) = _chunk_as_variate(z, mspace_flatsize(μ))
+@inline _chunk_as_variate(z::AbstractVector, ::Tuple{}) = z[begin]
+@inline _chunk_as_variate(z::AbstractVector, ::Any) = z
+
+
+"""
+    transport_to(ν, ::Type{MU}) where {MU<:StdMeasure}
+    transport_to(::Type{NU}, μ) where {NU<:StdMeasure}
+
+As a user convenience, a standard measure type like [`StdUniform`](@ref),
+[`StdExponential`](@ref), [`StdNormal`](@ref) or [`StdLogistic`](@ref)
+may be used directly as the source or target of a measure transport.
+
+The transport partner is an instance of the standard measure for measures
+with scalar variates, and a power of it with
+[`MeasureBase.some_dof(μ)`](@ref) (resp. `ν`) elements otherwise.
+"""
+function transport_to(ν, ::Type{MU}) where {MU<:StdMeasure}
+    transport_to(ν, _std_tp_partner(MU, ν))
+end
+
+function transport_to(::Type{NU}, μ) where {NU<:StdMeasure}
+    transport_to(_std_tp_partner(NU, μ), μ)
+end
+
+function transport_to(::Type{NU}, ::Type{MU}) where {NU<:StdMeasure,MU<:StdMeasure}
+    throw(
+        ArgumentError(
+            "Can't construct a transport function between the types of two standard measures, need a measure instance on one side",
+        ),
+    )
+end
+
+function _std_tp_partner(::Type{M}, μ) where {M<:StdMeasure}
+    m = asmeasure(μ)
+    _std_tp_partner_bysize(M, mspace_flatsize(m), m)
+end
+_std_tp_partner_bysize(::Type{M}, ::Tuple{}, μ) where {M<:StdMeasure} = M()
+_std_tp_partner_bysize(::Type{M}, ::Any, μ) where {M<:StdMeasure} = M()^some_dof(μ)
+
+
+# Element-wise transport kernels for broadcasts and maps:
+struct _ToStd{S} <: Function end
+@inline (::_ToStd{S})(μ, x) where {S} = transport_to_std(S, μ, x)
+struct _FromStd{S} <: Function end
+@inline (::_FromStd{S})(μ, z) where {S} = transport_from_std(S, μ, z)
+
+# Flat vector of standard variates from an array of standard variates:
+@inline _flat_std_of(A::AbstractArray{<:Number}) = vec(A)
+@inline _flat_std_of(A::AbstractArray{<:AbstractVector}) = _flatten_to_rv(vec(A))
+@inline _flat_std_of(A::AbstractVector{<:AbstractVector}) = _flatten_to_rv(A)
diff --git a/src/utils.jl b/src/utils.jl
index f73de9a4..cd215682 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -183,8 +183,12 @@ convert_realtype(::Type{T}, x::Union{Tuple,NamedTuple}) where {T<:Real} =
 convert_realtype(::Type{T}, x::AbstractArray) where {T<:Real} =
     map(Base.Fix1(convert_realtype, T), x)
 
-# Distributions implementation hooks:
-function _trafo_cdf_impl end
+# Distributions implementation hooks, specialized for dual numbers in the
+# ForwardDiff extension:
+function _trafo_logcdf_impl end
+function _trafo_logccdf_impl end
 function _trafo_quantile_impl end
-function _trafo_quantile_impl_generic end
+function _trafo_cquantile_impl end
+function _dist_quantile end
+function _dist_cquantile end
 function _dist_params_numtype end
diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl
index efd2cabd..413f2890 100644
--- a/test/combinators/bind.jl
+++ b/test/combinators/bind.jl
@@ -65,14 +65,14 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca
         xy_reco = transport_to(μ, StdUniform()^3)(y)
         @test xy_reco ≈ xy
 
-        # Transport between two measures of unknown DOF (mvstd pivot):
+        # Transport between two measures of unknown DOF (standard pivot):
         μ2 = mbind(f_βv, StdUniform()^1, vcat)
         xy2 = transport_to(μ2, μ)(xy)
         @test xy2 isa AbstractVector{<:Real} && length(xy2) == 3
         @test transport_to(μ, μ2)(xy2) ≈ xy
         @test logdensityof(μ2, xy2) isa Real
 
-        # Transport between known-DOF and unknown-DOF measures (mvstd pivot):
+        # Transport between known-DOF and unknown-DOF measures (standard pivot):
         ν_known = productmeasure((StdNormal(), StdNormal(), StdNormal()))
         z = transport_to(ν_known, μ)(xy)
         @test z isa Tuple{Vararg{Real,3}}
@@ -107,7 +107,7 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca
         yP = rand(stblrng(), Float64, P)
         z = transport_to(StdUniform()^6, P)(yP)
         @test z isa AbstractVector{<:Real} && length(z) == 6
-        yP_reco, rest = MeasureBase.transport_from_mvstd_with_rest(P, StdUniform(), z)
+        yP_reco, rest = MeasureBase.transport_from_std_with_rest(StdUniform, P, z)
         @test yP_reco isa Vector{<:AbstractVector{Float64}}
         @test yP_reco ≈ yP && isempty(rest)
         @test logdensityof(P, yP) ≈ logdensityof(μ, yP[1]) + logdensityof(μ, yP[2])
diff --git a/test/distributions/test_mooncake.jl b/test/distributions/test_mooncake.jl
index 4a20e46c..17c08427 100644
--- a/test/distributions/test_mooncake.jl
+++ b/test/distributions/test_mooncake.jl
@@ -48,8 +48,8 @@ _test_gradient(f, x::AbstractVector) = @test _mooncake_gradient(f, x) ≈ Forwar
         _test_gradient(x -> sum(transport_to(StdNormal()^3, asmeasure(pd))(x)), [0.4, 0.8, 1.5])
 
         dirich = Dirichlet([2.0, 3.0, 4.0])
-        _test_gradient(u -> MeasureBase.from_origin(dirich, u)[1], [0.3, 0.7])
-        _test_gradient(x -> sum(MeasureBase.to_origin(dirich, vcat(x, 1 - sum(x)))), [0.28, 0.23])
+        _test_gradient(u -> MeasureBase.transport_from_std(StdUniform, dirich, u)[1], [0.3, 0.7])
+        _test_gradient(x -> sum(MeasureBase.transport_to_std(StdUniform, dirich, vcat(x, 1 - sum(x)))), [0.28, 0.23])
     end
 
     @testset "logdensityof gradients" begin
diff --git a/test/distributions/test_shape_contract.jl b/test/distributions/test_shape_contract.jl
index b4f11ad7..52d5b608 100644
--- a/test/distributions/test_shape_contract.jl
+++ b/test/distributions/test_shape_contract.jl
@@ -19,8 +19,8 @@ using LinearAlgebra: I
     @test @inferred(preferred_stdmeasure(Uniform(1, 2))) === StdUniform
     @test @inferred(preferred_stdmeasure(Exponential(2.0))) === StdExponential
     @test @inferred(preferred_stdmeasure(Logistic(1, 2))) === StdLogistic
-    @test @inferred(preferred_stdmeasure(Beta(2, 3))) === StdUniform
-    @test @inferred(preferred_stdmeasure(truncated(Normal(), 0, 1))) === StdUniform
+    @test @inferred(preferred_stdmeasure(Beta(2, 3))) === StdLogistic
+    @test @inferred(preferred_stdmeasure(truncated(Normal(), 0, 1))) === StdLogistic
     @test @inferred(preferred_stdmeasure(MvNormal(zeros(2), I(2)))) === StdNormal
     @test @inferred(preferred_stdmeasure(Dirichlet([1.0, 2.0]))) === StdUniform
     @test @inferred(preferred_stdmeasure(Poisson(3))) <: NoStdTransport
@@ -28,10 +28,10 @@ using LinearAlgebra: I
     @test @inferred(preferred_stdmeasure(StandardDist{Uniform}())) === StdUniform
 
     @test @inferred(preferred_stdmeasure(productmeasure((asmeasure(Beta(2, 3)), asmeasure(Normal()))))) === StdNormal
-    @test @inferred(preferred_stdmeasure(productmeasure((a = Dirac(1.0), b = asmeasure(Beta(2, 3)))))) === StdUniform
+    @test @inferred(preferred_stdmeasure(productmeasure((a = Dirac(1.0), b = asmeasure(Beta(2, 3)))))) === StdLogistic
     @test @inferred(preferred_stdmeasure(productmeasure((a = asmeasure(Poisson(2)), b = asmeasure(Beta(2, 3)))))) <: NoStdTransport
     @test @inferred(preferred_stdmeasure(productmeasure([asmeasure(Normal(i, 1)) for i in 1:3]))) === StdNormal
-    @test @inferred(preferred_stdmeasure(product_distribution([Beta(2, 3), Beta(1, 1)]))) === StdUniform
+    @test @inferred(preferred_stdmeasure(product_distribution([Beta(2, 3), Beta(1, 1)]))) === StdLogistic
 
     lkj = asmeasure(LKJCholesky(3, 1.0))
     @test @inferred(mspace_elsize(lkj)) isa MeasureBase.NoMSpaceElementSize
diff --git a/test/distributions/test_transport.jl b/test/distributions/test_transport.jl
index f4692124..f7515e18 100644
--- a/test/distributions/test_transport.jl
+++ b/test/distributions/test_transport.jl
@@ -6,12 +6,13 @@ using LinearAlgebra
 using InverseFunctions, ChangesOfVariables
 using Distributions, ArraysOfArrays
 using StableRNGs
+using LogExpFunctions: logit
 import ForwardDiff, Zygote
 import PDMats
 
-using MeasureBase: transport_to, transport_def, transport_origin
-using MeasureBase: StdUniform, StdNormal, StdExponential
-using .MeasureBaseDistributionsExt: _trafo_cdf, _trafo_quantile
+using MeasureBase: transport_to, transport_def
+using MeasureBase: StdUniform, StdNormal, StdExponential, StdLogistic
+using .MeasureBaseDistributionsExt: _trafo_logcdf, _trafo_logccdf, _trafo_quantile, _trafo_cquantile
 
 include("getjacobian.jl")
 
@@ -20,12 +21,12 @@ include("getjacobian.jl")
     function test_back_and_forth(trg, src)
         @testset "transform $(typeof(trg).name) <-> $(typeof(src).name)" begin
             x = rand(src)
-            y = transport_def(trg, src, x)
-            src_v_reco = transport_def(src, trg, y)
+            y = transport_to(trg, src)(x)
+            src_v_reco = transport_to(src, trg)(y)
 
             @test x ≈ src_v_reco
-            
-            f = x -> transport_def(trg, src, x)
+
+            f = x -> transport_to(trg, src)(x)
             ref_ladj = logpdf(src, x) - logpdf(trg, y)
             @test ref_ladj ≈ logabsdet(getjacobian(f, x))[1]
         end
@@ -115,12 +116,47 @@ include("getjacobian.jl")
 
     @testset "Custom cdf and quantile for dual numbers" begin
         Dual = ForwardDiff.Dual
+        dual_normal = Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0))
+        dual_x = Dual(0.5, 0, 0, 1)
+        dual_p = Dual(0.3, 0, 0, 1)
+
+        @test isapprox(_trafo_logcdf(dual_normal, dual_x), logcdf(dual_normal, dual_x), rtol = 10^-6)
+        @test isapprox(_trafo_logcdf(Normal(0, 1), Dual(0.5, 1)), logcdf(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6)
+        @test isapprox(_trafo_logccdf(dual_normal, dual_x), logccdf(dual_normal, dual_x), rtol = 10^-6)
+        @test isapprox(_trafo_logccdf(Normal(0, 1), Dual(0.5, 1)), logccdf(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6)
+
+        @test isapprox(_trafo_quantile(Normal(0, 1), Dual(0.3, 1)), quantile(Normal(0, 1), Dual(0.3, 1)), rtol = 10^-6)
+        @test isapprox(_trafo_quantile(dual_normal, dual_p), quantile(dual_normal, dual_p), rtol = 10^-6)
+        @test isapprox(_trafo_cquantile(Normal(0, 1), Dual(0.3, 1)), cquantile(Normal(0, 1), Dual(0.3, 1)), rtol = 10^-6)
+        @test isapprox(_trafo_cquantile(dual_normal, dual_p), cquantile(dual_normal, dual_p), rtol = 10^-6)
+
+        # Distributions whose cdf doesn't support dual numbers natively:
+        beta = Beta(2.0, 3.0)
+        dlogitcdf(d, x) = pdf(d, x) / (cdf(d, x) * ccdf(d, x))
+        @test ForwardDiff.derivative(x -> transport_to(StdLogistic(), beta)(x), 0.3) ≈ dlogitcdf(beta, 0.3)
+        x_b = transport_to(beta, StdLogistic())(-0.4)
+        @test ForwardDiff.derivative(l -> transport_to(beta, StdLogistic())(l), -0.4) ≈ inv(dlogitcdf(beta, x_b))
+    end
 
-        @test isapprox(_trafo_cdf(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), cdf(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), rtol = 10^-6)
-        @test isapprox(_trafo_cdf(Normal(0, 1), Dual(0.5, 1)), cdf(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6)
-
-        @test isapprox(_trafo_quantile(Normal(0, 1), Dual(0.5, 1)), quantile(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6)
-        @test isapprox(_trafo_quantile(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), quantile(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), rtol = 10^-6)
+    @testset "tails of univariate transports" begin
+        # Bounded and heavy-lower-tailed distributions lose the lower tail in
+        # their quantile functions, so the ranges differ:
+        for (d, ls) in [
+            (Normal(0.3, 1.7), [-700.0, -40.0, -8.0, 0.0, 8.0, 40.0, 700.0]),
+            (Weibull(0.7, 1.3), [-40.0, -8.0, 0.0, 8.0, 40.0, 700.0]),
+            (truncated(Normal(0.2, 1.1), -3.0, 2.5), [-8.0, 0.0, 8.0]),
+        ]
+            for l in ls
+                x = transport_to(d, StdLogistic())(l)
+                @test insupport(d, x)
+                @test isapprox(transport_to(StdLogistic(), d)(x), l, rtol = 1e-6, atol = 1e-12)
+            end
+        end
+        for z in [-8.0, 8.0, 37.0]
+            x = transport_to(Weibull(0.7, 1.3), StdNormal())(z)
+            @test isfinite(x) && x > 0
+            @test transport_to(StdNormal(), Weibull(0.7, 1.3))(x) ≈ z rtol = 1e-6
+        end
     end
 
     @testset "trafo autodiff pullbacks" begin
diff --git a/test/shape_contract.jl b/test/shape_contract.jl
index 7c1ebc20..6e34973b 100644
--- a/test/shape_contract.jl
+++ b/test/shape_contract.jl
@@ -69,7 +69,7 @@ struct _CustomStd <: MeasureBase.StdMeasure end
         end
         # Transports of the base measure don't transport restricted measures:
         @test @inferred(preferred_stdmeasure(restrict(x -> x > 0, StdNormal()))) <: NoStdTransport
-        @test @inferred(preferred_stdmeasure(MeasureBase.Half(StdNormal()))) <: NoStdTransport
+        @test @inferred(preferred_stdmeasure(MeasureBase.Half(StdNormal()))) === StdUniform
 
         @test @inferred(preferred_stdmeasure(Dirac(2.0))) === AnyStdMeasure
         @test @inferred(preferred_stdmeasure(Lebesgue())) <: NoStdTransport
diff --git a/test/test_mooncake.jl b/test/test_mooncake.jl
index a6582c1d..7d10fb2c 100644
--- a/test/test_mooncake.jl
+++ b/test/test_mooncake.jl
@@ -9,7 +9,7 @@ import ForwardDiff
 using MeasureBase
 using MeasureBase: transport_to
 using MeasureBase: isneginf, isposinf, _adignore_call
-using MeasureBase: check_dof, require_insupport, _origin_depth
+using MeasureBase: check_dof, require_insupport
 
 _mooncake_gradient(f, x) = Mooncake.value_and_gradient!!(
     Mooncake.prepare_gradient_cache(f, x), f, x
@@ -25,7 +25,6 @@ _mooncake_gradient(f, x) = Mooncake.value_and_gradient!!(
         Mooncake.TestUtils.test_rule(rng, _adignore_call, () -> 42.0; is_primitive = true)
         Mooncake.TestUtils.test_rule(rng, check_dof, StdNormal(), StdUniform(); is_primitive = true)
         Mooncake.TestUtils.test_rule(rng, require_insupport, StdNormal(), 0.5; is_primitive = true)
-        Mooncake.TestUtils.test_rule(rng, _origin_depth, StdNormal(); is_primitive = true)
     end
 
     @testset "@_adignore is ignored" begin
diff --git a/test/transport.jl b/test/transport.jl
index 34359784..5a68d6b1 100644
--- a/test/transport.jl
+++ b/test/transport.jl
@@ -2,18 +2,14 @@ using Test
 
 using MeasureBase.Interface: transport_to, test_transport
 using MeasureBase: StdUniform, StdExponential, StdLogistic, StdNormal
-using MeasureBase: Dirac
+using MeasureBase: Dirac, Half, restrict, mbind, productmeasure, pushfwd
+using MeasureBase: transport_to_std, transport_from_std, transport_from_std_with_rest
+using InverseFunctions: inverse
+using StaticArrays: SVector
+using Static: static
 using LogExpFunctions: logit
 
-using ChainRulesTestUtils
-
 @testset "transport_to" begin
-    test_rrule(
-        MeasureBase._origin_depth,
-        pushfwd(exp, StdUniform()),
-        output_tangent = static(0),
-    )
-
     for (f, μ) in [
         (logit, StdUniform())
         (log, StdExponential())
@@ -60,6 +56,83 @@ using ChainRulesTestUtils
               transport_to(StdUniform()^(2, 3), StdExponential()^6)
     end
 
+    # Tail accuracy of transports between standard measures:
+
+    @testset "transports between standard measures" begin
+        stds = (StdUniform(), StdExponential(), StdLogistic(), StdNormal())
+        for ν in stds, μ in stds
+            f = transport_to(ν, μ)
+            for x in [rand(μ) for _ in 1:5]
+                @test inverse(f)(f(x)) ≈ x
+            end
+        end
+
+        # Round trips between the unbounded standard measures keep the tails:
+        for z in (-37.0, -20.0, -8.0, -6.0, 6.0, 8.0, 20.0, 37.0)
+            for ν in (StdExponential(), StdLogistic())
+                y = transport_to(ν, StdNormal())(z)
+                @test isfinite(y)
+                @test transport_to(StdNormal(), ν)(y) ≈ z rtol = 1e-8
+            end
+        end
+        for l in (-700.0, -40.0, -8.0, 8.0, 40.0, 700.0)
+            y = transport_to(StdExponential(), StdLogistic())(l)
+            @test isfinite(y) && y >= 0
+            @test transport_to(StdLogistic(), StdExponential())(y) ≈ l rtol = 1e-8
+        end
+        # The lower tail survives a uniform pivot, the upper tail saturates:
+        @test transport_to(StdNormal(), StdUniform())(transport_to(StdUniform(), StdNormal())(-37.0)) ≈ -37.0 rtol = 1e-8
+    end
+
+    @testset "scalar and static transports" begin
+        f = transport_to(StdNormal(), StdUniform())
+        @test @inferred(f(0.3)) isa Float64
+        @test @allocated(f(0.3)) == 0
+        g = transport_to(StdExponential()^static(3), StdNormal()^static(3))
+        xs = SVector(0.1, -0.4, 2.0)
+        @test @inferred(g(xs)) isa SVector{3,Float64}
+        @test @allocated(g(xs)) == 0
+        @test inverse(g)(g(xs)) ≈ xs
+        h = transport_to(StdNormal()^3, StdUniform()^3)
+        @test h(Float32[0.1, 0.5, 0.9]) isa Vector{Float32}
+    end
+
+    @testset "nested powers" begin
+        μ = (StdNormal()^2)^3
+        x = rand(μ)
+        f = transport_to(StdUniform()^6, μ)
+        y = f(x)
+        @test y isa AbstractVector{<:Real} && length(y) == 6
+        x_reco = inverse(f)(y)
+        @test all(map(≈, x_reco, x))
+        test_transport(StdExponential()^(3, 2), μ)
+    end
+
+    @testset "powers of measures without fast DOF" begin
+        f_β(a) = StdNormal()^length(a)
+        μ = mbind(f_β, StdUniform()^1, vcat)
+        P = μ^2
+        x = [rand(μ), rand(μ)]
+        z = transport_to(StdUniform()^4, P)(x)
+        @test z isa AbstractVector{<:Real} && length(z) == 4
+        x_reco = transport_to(P, StdUniform()^4)(z)
+        @test x_reco isa AbstractVector && all(map(≈, x_reco, x))
+    end
+
+    @testset "Half" begin
+        μ = Half(StdNormal())
+        test_transport(StdUniform(), μ)
+        test_transport(StdLogistic(), μ)
+        test_transport(μ, StdNormal())
+        @test transport_to(StdUniform(), μ)(0.0) ≈ 0
+    end
+
+    @testset "measures without standard transport" begin
+        μ = restrict(x -> x > 0, StdNormal())
+        @test_throws ArgumentError transport_to(StdUniform(), μ)(0.5)
+        @test_throws ArgumentError transport_to(μ, StdUniform())(0.5)
+    end
+
     @testset "transport for products" begin
         test_transport(
             StdUniform()^(2, 2),

From ffb6e719bf03ed48a741a16f28f86711914d8566 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 00:56:37 +0200
Subject: [PATCH 087/122] Fix batched kernels of static and array-variate
 products

Static flat sizes now work in variate streams and batched stream
consumption. Products over arrays of marginals with array variates
evaluate their marginals over slices of the flat batch instead of
handing flat storage to the nested point kernel, and their point
density accepts flat storage. The broadcast kernel of array products is
limited to concrete marginal types, product densities check the variate
structure instead of zipping silently, mcombine only merges homogeneous
products of concrete type, vcat-combined batches add lazily, and powers
with static axes report their flat size at the type level.

Created by generative AI.
---
 src/collection_utils.jl      |  3 ++
 src/combinators/combined.jl  | 13 +++---
 src/combinators/power.jl     |  6 +++
 src/combinators/product.jl   | 91 +++++++++++++++++++++++++++++++++---
 src/density-batched.jl       | 45 +++++++++++-------
 test/combinators/combined.jl | 23 ++++++++-
 test/logdensities.jl         | 74 +++++++++++++++++++++++++++++
 test/reactant/runtests.jl    | 21 +++++++++
 8 files changed, 246 insertions(+), 30 deletions(-)

diff --git a/src/collection_utils.jl b/src/collection_utils.jl
index 8ab2274a..009bfe00 100644
--- a/src/collection_utils.jl
+++ b/src/collection_utils.jl
@@ -109,6 +109,9 @@ function _consume_from_stream(x::AbstractVector, sz::Tuple{Vararg{IntegerLike}})
     return maybestatic_reshape(a_flat, sz), x_rest
 end
 
+Base.@propagate_inbounds _consume_from_stream(x::AbstractVector, sz::StaticArrays.Size) =
+    _consume_from_stream(x, _size_dims(sz))
+
 function _consume_from_stream(x::AbstractVector, @nospecialize(sz))
     throw(ArgumentError("Can't consume a variate of size $sz from a flat vector stream"))
 end
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index 886611f5..b1a211f1 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -84,8 +84,9 @@ end
     _mcombine_product_shortcut(f_c, marginals(α), marginals(β), α, β)
 end
 
-_mcombine_product_shortcut(::typeof(vcat), ma::AbstractVector{T}, mb::AbstractVector{T}, α, β) where {T} =
-    productmeasure(vcat(ma, mb))
+function _mcombine_product_shortcut(::typeof(vcat), ma::AbstractVector{T}, mb::AbstractVector{T}, α, β) where {T}
+    isconcretetype(T) ? productmeasure(vcat(ma, mb)) : _generic_mcombine_impl_stage2(vcat, α, β)
+end
 _mcombine_product_shortcut(::typeof(merge), ma::NamedTuple, mb::NamedTuple, α, β) =
     productmeasure(merge(ma, mb))
 _mcombine_product_shortcut(f_c, ma, mb, α, β) = _generic_mcombine_impl_stage2(f_c, α, β)
@@ -129,9 +130,9 @@ end
     _vcat_flatsize(mspace_flatsize(μ.α), mspace_flatsize(μ.β))
 end
 
-@inline _vcat_flatsize(a::SizeLike, b::SizeLike) = (size2length(a) + size2length(b),)
-@inline _vcat_flatsize(a::NoMSpaceElementSize, ::Any) = a
-@inline _vcat_flatsize(::Any, b::NoMSpaceElementSize) = b
+@inline _vcat_flatsize(a::SizeLike, b::SizeLike) = canonical_size((size2length(a) + size2length(b),))
+@inline _vcat_flatsize(a::NoMSpaceElementSize, ::SizeLike) = a
+@inline _vcat_flatsize(::SizeLike, b::NoMSpaceElementSize) = b
 @inline _vcat_flatsize(a::NoMSpaceElementSize, ::NoMSpaceElementSize) = a
 
 @inline getdof(μ::CombinedMeasure) = getdof(μ.α) + getdof(μ.β)
@@ -198,7 +199,7 @@ function batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, A::Ab
     ℓ_b, _, A_rest = batched_logdensityof_with_rest(μ.β, A2)
     n_μ = size(A, 1) - size(A_rest, 1)
     A_μ = view(A, 1:n_μ, Base.tail(axes(A))...)
-    return ℓ_a .+ ℓ_b, A_μ, A_rest
+    return _lazy_add(ℓ_a, ℓ_b), A_μ, A_rest
 end
 
 function logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector)
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 80bc1c2c..a33b56b1 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -79,6 +79,12 @@ marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes)
 
 @inline mspace_elsize(μ::PowerMeasure) = pwr_size(μ)
 @inline mspace_flatsize(μ::PowerMeasure) = _cat_sizes(mspace_flatsize(pwr_base(μ)), pwr_size(μ))
+@inline function mspace_flatsize(::Type{<:PowerMeasure{M,A}}) where {M,A<:Tuple{Vararg{StaticOneToLike}}}
+    _cat_sizes(mspace_flatsize(M), _static_axes_size(A))
+end
+@generated function _static_axes_size(::Type{A}) where {A<:Tuple{Vararg{StaticOneToLike}}}
+    :(StaticArrays.Size($(map(T -> T.parameters[1], A.parameters)...)))
+end
 
 function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{AbstractArray,N}}) where {N}
     powermeasure(μ, dims)
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index f78fddbc..0f92f570 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -74,10 +74,30 @@ end
 
 for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)]
     @eval @inline function $head(d::AbstractProductMeasure, x)
+        _check_marginal_count(marginals(d), x)
         mapreduce($func, +, marginals(d), x)
     end
 end
 
+# Variates of products are collections of marginal variates, with the same
+# structure as the marginals:
+@inline function _check_marginal_count(mar::AbstractArray, x::AbstractArray)
+    size(mar) == size(x) || _throw_marginal_mismatch()
+    return nothing
+end
+@inline _check_marginal_count(::AbstractArray, x) = _throw_marginal_mismatch()
+# Tuple products also take vector variates (e.g. from converted product
+# distributions):
+@inline function _check_marginal_count(mar::Tuple, x::Union{Tuple,AbstractVector})
+    length(mar) == length(x) || _throw_marginal_mismatch()
+    return nothing
+end
+@inline _check_marginal_count(::Tuple, x) = _throw_marginal_mismatch()
+@inline _check_marginal_count(mar, x) = nothing
+
+@noinline _throw_marginal_mismatch() =
+    throw(ArgumentError("Variate doesn't match the structure of the marginals of a product measure"))
+
 struct ProductMeasure{M} <: AbstractProductMeasure
     marginals::M
 end
@@ -129,6 +149,7 @@ end
 for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)]
     # For tuples, `mapreduce` has trouble with type inference
     @eval @inline function $head(d::ProductMeasure{T}, x) where {T<:Tuple}
+        _check_marginal_count(marginals(d), x)
         ℓs = map($func, marginals(d), x)
         sum(ℓs)
     end
@@ -205,21 +226,57 @@ marginals(μ::ProductMeasure) = μ.marginals
     _cat_sizes(mspace_flatsize(M), maybestatic_size(marginals(μ)))
 end
 
-# The marginals align with the leading dimensions of the flat batch, so
-# one broadcast evaluates all marginal densities:
+# Batched densities over flat storage `(marginal flat dims..., product
+# dims..., batch dims...)`. Marginals with scalar variates align with the
+# leading dimensions of the batch, so one broadcast evaluates all marginal
+# densities. Marginals with array variates are evaluated one by one over
+# their slices of the batch.
 @inline function batched_logdensityof_impl(μ::ProductMeasure{<:AbstractArray{M,N}}, A::AbstractArray) where {M,N}
-    _product_batched_ld(μ, A, mspace_flatsize(M), Val(N))
+    _product_batched_ld(μ, A, mspace_flatsize(M), Val(N), Val(isconcretetype(M)))
 end
 
-@inline function _product_batched_ld(μ::ProductMeasure, A::AbstractArray, ::Tuple{}, ::Val{N}) where {N}
+@inline function _product_batched_ld(μ::ProductMeasure, A::AbstractArray, ::Tuple{}, ::Val{N}, ::Val{true}) where {N}
     ld = Broadcast.instantiate(Broadcast.broadcasted(dynamic ∘ logdensityof_impl, marginals(μ), A))
     _sum_leading_dims(ld, static(N))
 end
 
-@inline function _product_batched_ld(μ::ProductMeasure, A::AbstractArray, ::Any, ::Val)
+@inline function _product_batched_ld(μ::ProductMeasure, A::AbstractArray, sz_m::SizeLike, ::Val{N}, ::Val{true}) where {N}
+    _marginal_slices_ld(marginals(μ), A, Val(length(sz_m)), Val(ndims(A) - length(sz_m) - N))
+end
+
+@inline function _product_batched_ld(μ::ProductMeasure, A::AbstractArray, ::Any, ::Val, ::Val)
     _batched_ld_generic(logdensityof_impl, μ, A)
 end
 
+function _marginal_slices_ld(mar::AbstractArray{<:Any,N}, A::AbstractArray, ::Val{K}, ::Val{B}) where {N,K,B}
+    if ndims(A) != K + N + B || ntuple(i -> size(A, K + i), Val(N)) != size(mar)
+        _throw_size_mismatch()
+    end
+    lead = ntuple(_ -> Colon(), Val(K))
+    trail = ntuple(_ -> Colon(), Val(B))
+    ld(i) = _materialize(batched_logdensityof_impl(mar[i], view(A, lead..., Tuple(i)..., trail...)))
+    init = _zero_logd(A, ntuple(i -> size(A, K + N + i), Val(B)))
+    return mapreduce(ld, +, CartesianIndices(mar); init = init)
+end
+
+@inline _zero_logd(A::AbstractArray, ::Tuple{}) = zero(_logd_numtype(A))
+@inline _zero_logd(A::AbstractArray, dims::Tuple) = fill!(similar(A, _logd_numtype(A), dims), 0)
+
+# The point density of array products with array-variate marginals accepts
+# the flat variate storage `(marginal flat dims..., product dims...)`:
+@inline function logdensityof_impl(μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray{<:Number}) where {M}
+    _array_product_ld(μ, x, mspace_flatsize(M))
+end
+@inline _array_product_ld(μ::ProductMeasure, x::AbstractArray, ::Tuple{}) = _array_product_ld_nested(μ, x)
+@inline _array_product_ld(μ::ProductMeasure, x::AbstractArray, ::NoMSpaceElementSize) = _array_product_ld_nested(μ, x)
+@inline function _array_product_ld(μ::ProductMeasure, x::AbstractArray, sz_m::SizeLike)
+    _marginal_slices_ld(marginals(μ), x, Val(length(sz_m)), Val(0))
+end
+@inline function _array_product_ld_nested(μ::ProductMeasure, x::AbstractArray)
+    _check_marginal_count(marginals(μ), x)
+    mapreduce(logdensityof, +, marginals(μ), x)
+end
+
 # TODO: Better `map` support in MappedArrays
 _map(f, args...) = map(f, args...)
 _map(f, x::MappedArrays.ReadonlyMappedArray) = mappedarray(fchain((x.f, f)), x.data)
@@ -277,10 +334,30 @@ function checked_arg(μ::ProductMeasure{<:NTuple{N,Any}}, x::NTuple{N,Any}) wher
     map(checked_arg, marginals(μ), x)
 end
 
-function checked_arg(μ::ProductMeasure{<:AbstractArray}, x::AbstractArray)
-    map(checked_arg, marginals(μ), x)
+# Variates of array products are arrays of marginal variates or, for
+# marginals with array variates of known size, their flat storage:
+@propagate_inbounds function checked_arg(μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray) where {M}
+    @boundscheck _check_product_arg(marginals(μ), x, mspace_flatsize(M))
+    return x
 end
 
+@inline _check_product_arg(mar, x::AbstractArray, ::Tuple{}) = _check_marginal_count(mar, x)
+@inline _check_product_arg(mar, x::AbstractArray{<:Number}, ::NoMSpaceElementSize) = _check_marginal_count(mar, x)
+@inline function _check_product_arg(mar, x::AbstractArray, ::NoMSpaceElementSize)
+    _check_marginal_count(mar, x)
+    foreach(checked_arg, mar, x)
+    return nothing
+end
+@inline function _check_product_arg(mar, x::AbstractArray, sz_m::SizeLike)
+    if size(x) == size(mar)
+        foreach(checked_arg, mar, x)
+    elseif size(x) != (Tuple(sz_m)..., size(mar)...)
+        _throw_marginal_mismatch()
+    end
+    return nothing
+end
+
+
 function checked_arg(
     μ::ProductMeasure{<:NamedTuple{names}},
     x::NamedTuple{names},
diff --git a/src/density-batched.jl b/src/density-batched.jl
index 26ddb3c3..9ad63765 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -249,8 +249,8 @@ end
 end
 
 
-@inline _lazy_add(c, x::Number) = c + x
-@inline _lazy_add(c, A) = Broadcast.instantiate(Broadcast.broadcasted(+, c, A))
+@inline _lazy_add(a::Number, b::Number) = a + b
+@inline _lazy_add(a, b) = Broadcast.instantiate(Broadcast.broadcasted(+, a, b))
 
 
 """
@@ -261,35 +261,48 @@ Batched form of [`MeasureBase.logdensityof_with_rest`](@ref) for a batch
 streams, all further dimensions are batch dimensions.
 
 Returns a tuple `(ℓ, A_μ, A_rest)` of the log-densities over the batch
-dimensions, the flat variate batch consumed from the streams and the
-unconsumed rest of the streams.
+dimensions (possibly as a lazy broadcast), the rows consumed from the
+streams and the unconsumed rest of the streams.
 
-Requires the flat variate size of `μ` to be known, see
-[`MeasureBase.mspace_flatsize`](@ref).
+Consuming from streams requires the flat variate sizes of `μ` or its
+components to be known, see [`MeasureBase.mspace_flatsize`](@ref).
 """
 function batched_logdensityof_with_rest end
 
 function batched_logdensityof_with_rest(μ::AbstractMeasure, A::AbstractArray)
     A_μ, A_rest = _batched_consume(A, mspace_flatsize(μ))
-    return _materialize(_batched_ld(logdensityof_impl, μ, A_μ)), A_μ, A_rest
+    return _batched_ld(logdensityof_impl, μ, A_μ), A_μ, A_rest
 end
 
 # Consume the leading rows of a batch of streams as a batch of flat variates:
 @inline function _batched_consume(A::AbstractArray, sz::SizeLike)
-    n = size2length(sz)
-    n_stream = size(A, 1)
-    if n_stream < n
+    A_flat, A_rest = _batched_split(A, dynamic(size2length(sz)))
+    return _batched_chunk_shape(A_flat, _size_dims(sz)), A_rest
+end
+
+@inline function _batched_consume(A::AbstractArray, ::Tuple{})
+    A_flat, A_rest = _batched_split(A, 1)
+    batch_axes = Base.tail(axes(A))
+    return view(A_flat, firstindex(A_flat, 1), batch_axes...), A_rest
+end
+
+@inline function _batched_split(A::AbstractArray, n::Integer)
+    stream_idxs = axes(A, 1)
+    if length(stream_idxs) < n
         throw(ArgumentError("Variate streams too short during batched density evaluation"))
     end
     batch_axes = Base.tail(axes(A))
-    A_flat = view(A, 1:dynamic(n), batch_axes...)
-    A_rest = view(A, (dynamic(n) + 1):n_stream, batch_axes...)
-    return maybestatic_reshape(A_flat, (_size_dims(sz)..., map(length, batch_axes)...)), A_rest
+    i0 = first(stream_idxs)
+    A_flat = view(A, i0:(i0 + n - 1), batch_axes...)
+    A_rest = view(A, (i0 + n):last(stream_idxs), batch_axes...)
+    return A_flat, A_rest
 end
 
-@inline function _batched_consume(A::AbstractArray, ::Tuple{})
-    batch_axes = Base.tail(axes(A))
-    view(A, 1, batch_axes...), view(A, 2:size(A, 1), batch_axes...)
+# Chunks of variates with more than one flat dimension are reshaped, the
+# batch dimensions are dynamic anyway:
+@inline _batched_chunk_shape(A_flat::AbstractArray, ::Tuple{IntegerLike}) = A_flat
+@inline function _batched_chunk_shape(A_flat::AbstractArray, dims::Tuple{Vararg{IntegerLike}})
+    reshape(A_flat, (map(dynamic, dims)..., Base.tail(size(A_flat))...))
 end
 
 function _batched_consume(::AbstractArray, sz::NoMSpaceElementSize)
diff --git a/test/combinators/combined.jl b/test/combinators/combined.jl
index 9a79f225..43d9e20c 100644
--- a/test/combinators/combined.jl
+++ b/test/combinators/combined.jl
@@ -8,7 +8,8 @@ using OneTwoMany: firstarg, secondarg
 
 using MeasureBase
 using MeasureBase: StdExponential, StdLogistic, StdNormal, StdUniform
-using MeasureBase: mcombine, productmeasure, transport_to
+using MeasureBase: mcombine, productmeasure, transport_to, pushfwd
+using AffineMaps: Mul
 
 @testset "mcombine" begin
     stblrng() = StableRNG(789990641)
@@ -32,6 +33,26 @@ using MeasureBase: mcombine, productmeasure, transport_to
               MeasureBase.Dirac((1, 2))
     end
 
+    @testset "mcombine of products" begin
+        p1 = productmeasure([pushfwd(Mul(1.0), StdNormal()), pushfwd(Mul(2.0), StdNormal())])
+        p2 = productmeasure([pushfwd(Mul(3.0), StdNormal())])
+        p12 = mcombine(vcat, p1, p2)
+        @test p12 isa MeasureBase.ProductMeasure && length(MeasureBase.marginals(p12)) == 3
+        x3 = randn(3)
+        @test logdensityof(p12, x3) ≈ logdensityof(p1, x3[1:2]) + logdensityof(p2, x3[3:3])
+
+        # Products of different or abstract marginal types stay combined measures:
+        mab = mcombine(vcat, StdNormal()^2, StdUniform()^2)
+        @test mab isa MeasureBase.CombinedMeasure
+        pa = productmeasure(AbstractMeasure[StdNormal(), StdNormal()])
+        pb = productmeasure(AbstractMeasure[StdUniform(), StdUniform()])
+        mpab = mcombine(vcat, pa, pb)
+        @test mpab isa MeasureBase.CombinedMeasure
+        x = vcat(randn(2), rand(2))
+        @test logdensityof(mpab, x) ≈ logdensityof(mab, x)
+        @test_throws ArgumentError logdensities(mpab, vcat(randn(2, 3), rand(2, 3)))
+    end
+
     @testset "CombinedMeasure" begin
         μ = mcombine(Pair, α, β)
         @test μ isa MeasureBase.CombinedMeasure
diff --git a/test/logdensities.jl b/test/logdensities.jl
index 98d1323b..5ecf21c4 100644
--- a/test/logdensities.jl
+++ b/test/logdensities.jl
@@ -13,6 +13,17 @@ using JLArrays: JLArray
 
 stdnormal_ld(x) = -(x^2 + log2π) / 2
 
+# A measure with array variates of a known flat size at the type level:
+struct VecTestMeasure{T} <: AbstractMeasure
+    s::T
+end
+MeasureBase.mspace_elsize(::VecTestMeasure) = (2,)
+MeasureBase.mspace_flatsize(::VecTestMeasure) = (2,)
+MeasureBase.mspace_flatsize(::Type{<:VecTestMeasure}) = (2,)
+MeasureBase.basemeasure(::VecTestMeasure) = LebesgueBase()^2
+MeasureBase.insupport(::VecTestMeasure, x) = true
+MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s)
+
 @testset "logdensities" begin
     @testset "scalar variates" begin
         X = randn(10)
@@ -187,6 +198,69 @@ stdnormal_ld(x) = -(x^2 + log2π) / 2
         @test @inferred(logdensities(m3, X3)) ≈ [logdensityof(m3, x) for x in eachcol(X3)]
     end
 
+    @testset "static sizes in combined measures" begin
+        m = mcombine(vcat, StdNormal()^static(2), StdUniform()^3)
+        @test @inferred(MeasureBase.mspace_flatsize(m)) == (5,)
+        x = vcat(randn(2), rand(3))
+        @test @inferred(logdensityof(m, x)) ≈ sum(stdnormal_ld, x[1:2])
+        X = vcat(randn(2, 4), rand(3, 4))
+        @test @inferred(logdensities(m, X)) ≈ [logdensityof(m, x) for x in eachcol(X)]
+        ms = mcombine(vcat, StdNormal()^static(2), StdUniform()^static(3))
+        @test @inferred(MeasureBase.mspace_flatsize(ms)) == MeasureBase.mspace_flatsize(StdNormal()^static(5))
+        @test @inferred(logdensityof(ms, SVector{5}(x))) ≈ logdensityof(m, x)
+        @test @inferred(logdensityof(ms, x)) ≈ logdensityof(m, x)
+        @test logdensities(ms, X) ≈ logdensities(m, X)
+        @test_throws ArgumentError MeasureBase.batched_logdensityof_with_rest(StdNormal(), zeros(0, 4))
+    end
+
+    @testset "array products of array-variate marginals" begin
+        p = productmeasure([VecTestMeasure(1.0), VecTestMeasure(2.0), VecTestMeasure(0.5)])
+        @test @inferred(MeasureBase.mspace_flatsize(p)) == (2, 3)
+        xs = [randn(2) for _ in 1:3]
+        X = stack(xs)
+        ℓ = sum(map(logdensityof, MeasureBase.marginals(p), xs))
+        @test @inferred(logdensityof(p, xs)) ≈ ℓ
+        @test @inferred(logdensityof(p, X)) ≈ ℓ
+        A = randn(2, 3, 5)
+        @test @inferred(logdensities(p, A)) ≈ [logdensityof(p, A[:, :, i]) for i in 1:5]
+        @test logdensities(p, sliced(A, Val(2))) ≈ logdensities(p, A)
+        @test logdensities(p, zeros(2, 3, 0)) == Float64[]
+        @test_throws ArgumentError logdensityof(p, randn(2, 2))
+        @test_throws ArgumentError logdensities(p, randn(2, 2, 5))
+
+        # Powers with static axes have a flat size at the type level:
+        pp = MeasureBase.ProductMeasure([StdNormal()^static(2), StdNormal()^static(2)])
+        @test @inferred(MeasureBase.mspace_flatsize(pp)) == (2, 2)
+        Xp = randn(2, 2)
+        @test @inferred(logdensityof(pp, Xp)) ≈ sum(stdnormal_ld, Xp)
+        A3 = randn(2, 2, 3)
+        @test @inferred(logdensities(pp, A3)) ≈ [sum(stdnormal_ld, A3[:, :, i]) for i in 1:3]
+    end
+
+    @testset "products with mixed marginal types" begin
+        pa = productmeasure(AbstractMeasure[StdNormal(), StdUniform()])
+        X = vcat(randn(1, 4), rand(1, 4))
+        xs = [X[:, i] for i in 1:4]
+        @test logdensities(pa, xs) ≈ [logdensityof(pa, x) for x in xs]
+        @test logdensities(pa, sliced(X, Val(1))) ≈ [logdensityof(pa, x) for x in xs]
+        @test_throws ArgumentError logdensityof(pa, 0.5)
+        @test_throws ArgumentError logdensityof(productmeasure((StdNormal(), StdUniform())), 0.5)
+        @test_throws ArgumentError logdensityof(pa, X[:, 1:1])
+    end
+
+    @testset "out-of-support and empty batches of structural kernels" begin
+        w = weightedmeasure(log(0.3), StdUniform()^2)
+        X = [0.5 -0.5 0.5; 0.5 0.5 1.5]
+        @test logdensities(w, X) == [log(0.3), -Inf, -Inf]
+        @test logdensities(w, zeros(2, 0)) == Float64[]
+        pu = productmeasure([weightedmeasure(log(i), StdUniform()) for i in 1:2])
+        @test logdensities(pu, X) == [log(2), -Inf, -Inf]
+        mc = mcombine(vcat, StdUniform()^1, StdExponential()^1)
+        @test logdensities(mc, [0.5 -0.5 0.5; 0.5 0.5 -1.0]) == [-0.5, -Inf, -Inf]
+        @test logdensities(mc, zeros(2, 0)) == Float64[]
+        @test MeasureBase.logdensity_def(pu, [0.5, 0.5]) ≈ log(2)
+    end
+
     @testset "GPU array semantics for structural kernels" begin
         JLArrays.allowscalar(false)
         ms = JLArray([weightedmeasure(log(i), StdNormal()) for i in 1:4])
diff --git a/test/reactant/runtests.jl b/test/reactant/runtests.jl
index f79e4754..182034b8 100644
--- a/test/reactant/runtests.jl
+++ b/test/reactant/runtests.jl
@@ -10,6 +10,7 @@ using Reactant
 using MeasureBase
 using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic, Lebesgue, Dirac
 using MeasureBase: logdensities, logdensity_rel, weightedmeasure, superpose, restrict, mintegrate_exp
+using MeasureBase: mcombine
 using ArraysOfArrays: VectorOfSimilarVectors, sliced
 using Distributions: Normal, Exponential, Uniform, Beta
 
@@ -67,6 +68,26 @@ _plain(x::Number) = Float64(x)
         test_traced(x -> logdensities(sm, x), vcat(x, 0.0))
     end
 
+    # Products over arrays of marginals are not covered: Reactant can't
+    # broadcast over arrays of measures together with traced arrays.
+    @testset "structural batched kernels" begin
+        w = weightedmeasure(log(0.3), StdNormal()^3)
+        test_traced(X -> logdensities(w, X), X)
+        mc = mcombine(vcat, StdNormal()^2, StdUniform()^1)
+        test_traced(X -> logdensities(mc, X), vcat(X[1:2, :], rand(1, 20)))
+        test_traced(X -> logdensities((StdNormal()^2)^3, X), reshape(X[1:2, 1:6], 2, 3, 2))
+    end
+
+    @testset "transport of powers and products" begin
+        test_traced(x -> transport_to(StdNormal()^10, StdUniform()^10)(x), rand(10))
+        test_traced(x -> transport_to(StdExponential()^10, StdNormal()^10)(x), x)
+        test_traced(x -> transport_to(StdLogistic()^10, StdExponential()^10)(x), rand(10))
+        test_traced(x -> transport_to(StdNormal()^6, (StdUniform()^2)^3)(x), rand(2, 3))
+        mc = mcombine(vcat, StdNormal()^2, StdUniform()^1)
+        test_traced(x -> transport_to(StdNormal()^3, mc)(x), vcat(randn(2), rand(1)))
+        test_traced(z -> transport_to(mc, StdNormal()^3)(z), randn(3))
+    end
+
     @testset "transport" begin
         test_traced(x -> transport_to(StdUniform(), StdNormal()).(x), x)
         test_traced(x -> transport_to(StdNormal(), StdUniform()).(x), rand(10))

From e86a3d3a615598e703e7ca9371ad3e5cabff05c0 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 01:06:08 +0200
Subject: [PATCH 088/122] Add batched transport and a broadcast hook for
 transport functions

Broadcasting a transport function over an array of variates with flat
storage (or over the flat storage of a batch) now transports the whole
batch in fused operations: batched_transport_to_std and
batched_transport_from_std work on flat batches with the variate
dimensions leading, streams of standard variates are consumed along
their first dimension. Powers, array products, weighted measures,
pushforwards of scalar-variate measures, vcat-combined measures and
MvNormal have batched kernels, everything else falls back to per-point
transport. Degrees of freedom of array products are summed dynamically.

Created by generative AI.
---
 .../MeasureBaseDistributionsExt.jl            |   1 +
 .../distribution_measure.jl                   |   4 +
 ext/MeasureBaseDistributionsExt/standardmv.jl |  12 ++
 src/MeasureBase.jl                            |   1 +
 src/combinators/combined.jl                   |  34 ++++
 src/combinators/power.jl                      |  53 ++++++
 src/combinators/product.jl                    |  36 +++-
 src/combinators/transformedmeasure.jl         |  22 +++
 src/combinators/weighted.jl                   |   9 +
 src/primitives/dirac.jl                       |  11 ++
 src/standard/stdconvert.jl                    |   6 +
 src/standard/stdmeasure.jl                    |   3 +
 src/transport-batched.jl                      | 156 ++++++++++++++++++
 test/distributions/test_transport.jl          |  17 ++
 test/reactant/runtests.jl                     |   9 +-
 test/transport.jl                             |  61 +++++++
 16 files changed, 432 insertions(+), 3 deletions(-)
 create mode 100644 src/transport-batched.jl

diff --git a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
index 79df0d97..586ae186 100644
--- a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
+++ b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
@@ -18,6 +18,7 @@ using MeasureBase: PowerMeasure, WeightedMeasure, SuperpositionMeasure, Pushforw
 using MeasureBase: basemeasure, rootmeasure, testvalue, productmeasure, pushfwd, superpose
 using MeasureBase: getdof, checked_arg, massof
 using MeasureBase: transport_to, transport_def, transport_to_std, transport_from_std
+using MeasureBase: batched_transport_to_std, batched_transport_from_std
 using MeasureBase: Reshape
 using MeasureBase: convert_realtype, _fwddiff, @_adignore
 import MeasureBase:
diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index d7a6eb6e..80c6f53e 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -79,6 +79,10 @@ end
     MeasureBase.transport_to_std(S, m.obj, x)
 @inline MeasureBase.transport_from_std(::Type{S}, m::DistributionMeasure, z) where {S<:StdMeasure} =
     MeasureBase.transport_from_std(S, m.obj, z)
+@inline MeasureBase.batched_transport_to_std(::Type{S}, m::DistributionMeasure, X::AbstractArray) where {S<:StdMeasure} =
+    MeasureBase.batched_transport_to_std(S, m.obj, X)
+@inline MeasureBase.batched_transport_from_std(::Type{S}, m::DistributionMeasure, Z::AbstractArray) where {S<:StdMeasure} =
+    MeasureBase.batched_transport_from_std(S, m.obj, Z)
 
 @inline MeasureBase.paramnames(m::DistributionMeasure) = propertynames(m.obj)
 @inline MeasureBase.params(m::DistributionMeasure) = NamedTuple{propertynames(m.obj)}(Distributions.params(m.obj))
diff --git a/ext/MeasureBaseDistributionsExt/standardmv.jl b/ext/MeasureBaseDistributionsExt/standardmv.jl
index b21a7ce5..1163c39b 100644
--- a/ext/MeasureBaseDistributionsExt/standardmv.jl
+++ b/ext/MeasureBaseDistributionsExt/standardmv.jl
@@ -19,6 +19,18 @@ function MeasureBase.transport_from_std(::Type{StdNormal}, d::MvNormal, z)
     muladd(_cholesky_L(d.Σ), z, d.μ)
 end
 
+function MeasureBase.batched_transport_to_std(::Type{StdNormal}, d::MvNormal, X::AbstractArray)
+    X_mat = reshape(X, (length(d), :))
+    Z_mat = _cholesky_L(d.Σ) \ (X_mat .- d.μ)
+    return reshape(Z_mat, size(X))
+end
+
+function MeasureBase.batched_transport_from_std(::Type{StdNormal}, d::MvNormal, Z::AbstractArray)
+    Z_mat = reshape(Z, (length(d), :))
+    X_mat = muladd(_cholesky_L(d.Σ), Z_mat, d.μ)
+    return reshape(X_mat, size(Z))
+end
+
 
 #DirichletMultinomial
 #Distributions.AbstractMvLogNormal
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 7c14fa92..7e62d783 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -208,6 +208,7 @@ include("combinators/superpose.jl")
 include("combinators/product.jl")
 include("combinators/power.jl")
 include("density-batched.jl")
+include("transport-batched.jl")
 include("combinators/spikemixture.jl")
 include("combinators/likelihood.jl")
 include("combinators/restricted.jl")
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index b1a211f1..7d9ecedb 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -265,3 +265,37 @@ function transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure, z::Abstrac
     b, z_rest = transport_from_std_with_rest(S, μ.β, z2)
     return μ.f_c(a, b), z_rest
 end
+
+
+# Batched transport consumes the variate parts of both component measures
+# along batches of streams:
+
+function batched_transport_to_std(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray) where {S<:StdMeasure}
+    Z, _, X_rest = batched_transport_to_std_with_rest(S, μ, X)
+    if size(X_rest, 1) != 0
+        throw(ArgumentError("Variate streams too long during batched transport of a combined measure"))
+    end
+    return Z
+end
+
+function batched_transport_to_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray) where {S<:StdMeasure}
+    Z_a, _, X2 = batched_transport_to_std_with_rest(S, μ.α, X)
+    Z_b, _, X_rest = batched_transport_to_std_with_rest(S, μ.β, X2)
+    X_μ, _ = _batched_split(X, size(X, 1) - size(X_rest, 1))
+    return vcat(Z_a, Z_b), X_μ, X_rest
+end
+
+function batched_transport_from_std(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, Z::AbstractArray) where {S<:StdMeasure}
+    X, Z_rest = batched_transport_from_std_with_rest(S, μ, Z)
+    if size(Z_rest, 1) != 0
+        throw(ArgumentError("Length of standard variates doesn't match degrees of freedom of a combined measure"))
+    end
+    return X
+end
+
+function batched_transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, Z::AbstractArray) where {S<:StdMeasure}
+    A, Z2 = batched_transport_from_std_with_rest(S, μ.α, Z)
+    B, Z_rest = batched_transport_from_std_with_rest(S, μ.β, Z2)
+    X = vcat(_as_stream_batch(A, mspace_flatsize(μ.α)), _as_stream_batch(B, mspace_flatsize(μ.β)))
+    return X, Z_rest
+end
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index a33b56b1..76afc6a5 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -265,3 +265,56 @@ end
 @inline _nest_leaf(A::AbstractArray, ::Val{M}, ::Val) where {M} = sliced(A, Val(M))
 @inline _pwr_nest(ν::PowerMeasure, B::AbstractArray) = sliced(B, Val(length(pwr_axes(ν))))
 @inline _pwr_nest(ν, B::AbstractArray) = B
+
+# Batched transport over the flat storage `(base variate dims..., power
+# dims..., batch dims...)`:
+
+function batched_transport_to_std(::Type{S}, μ::PowerMeasure, X::AbstractArray) where {S<:StdMeasure}
+    _pwr_batched_to_std(S, μ, X, mspace_flatsize(μ))
+end
+
+function _pwr_batched_to_std(::Type{S}, μ::PowerMeasure, X::AbstractArray, sz_flat::SizeLike) where {S}
+    ν, _ = _pwr_unwrap(μ)
+    _check_flatsize(X, sz_flat)
+    n_flat = length(sz_flat)
+    Z = _pwr_batched_to_std_flat(S, ν, X, mspace_flatsize(ν))
+    batch_dims = ntuple(i -> size(X, n_flat + i), Val(ndims(X) - n_flat))
+    return reshape(Z, (dynamic(fast_dof(μ)), batch_dims...))
+end
+
+function _pwr_batched_to_std(::Type{S}, μ::PowerMeasure, ::AbstractArray, ::NoMSpaceElementSize) where {S}
+    throw(ArgumentError("Batched transport of powers of measures of type $(nameof(typeof(pwr_base(μ)))) requires a known variate size"))
+end
+
+@inline function _pwr_batched_to_std_flat(::Type{S}, ν, X::AbstractArray, ::Tuple{}) where {S}
+    broadcast(Base.Fix1(_ToStd{S}(), ν), X)
+end
+
+@inline function _pwr_batched_to_std_flat(::Type{S}, ν, X::AbstractArray, sz::SizeLike) where {S}
+    stacked(map(Base.Fix1(_ToStd{S}(), ν), sliced(X, Val(length(sz)))))
+end
+
+function batched_transport_from_std(::Type{S}, μ::PowerMeasure, Z::AbstractArray) where {S<:StdMeasure}
+    _pwr_batched_from_std(S, μ, Z, mspace_flatsize(μ))
+end
+
+function _pwr_batched_from_std(::Type{S}, μ::PowerMeasure, Z::AbstractArray, sz_flat::SizeLike) where {S}
+    ν, _ = _pwr_unwrap(μ)
+    batch_dims = Base.tail(size(Z))
+    X = _pwr_batched_from_std_flat(S, ν, Z, batch_dims, mspace_flatsize(ν))
+    return reshape(X, (map(dynamic, _size_dims(sz_flat))..., batch_dims...))
+end
+
+function _pwr_batched_from_std(::Type{S}, μ::PowerMeasure, ::AbstractArray, ::NoMSpaceElementSize) where {S}
+    throw(ArgumentError("Batched transport to powers of measures of type $(nameof(typeof(pwr_base(μ)))) requires a known variate size"))
+end
+
+@inline function _pwr_batched_from_std_flat(::Type{S}, ν, Z::AbstractArray, batch_dims, ::Tuple{}) where {S}
+    broadcast(Base.Fix1(_FromStd{S}(), ν), Z)
+end
+
+@inline function _pwr_batched_from_std_flat(::Type{S}, ν, Z::AbstractArray, batch_dims, sz::SizeLike) where {S}
+    n_variates = size(Z, 1) ÷ dynamic(fast_dof(ν))
+    chunks = sliced(reshape(Z, (dynamic(fast_dof(ν)), n_variates, batch_dims...)), Val(1))
+    stacked(map(Base.Fix1(_FromStd{S}(), ν), chunks))
+end
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 0f92f570..b63821ca 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -327,8 +327,15 @@ end
 @inline _all_insupport(A::AbstractArray{<:NoFastInsupport{T}}) where {T} = NoFastInsupport{T}()
 @inline _all_insupport(A::AbstractArray) = all(A)
 
-getdof(d::AbstractProductMeasure) = sum(getdof, marginals(d))
-fast_dof(d::AbstractProductMeasure) = sum(fast_dof, marginals(d))
+getdof(d::AbstractProductMeasure) = _sum_dofs(getdof, marginals(d))
+fast_dof(d::AbstractProductMeasure) = _sum_dofs(fast_dof, marginals(d))
+
+# Sums over static DOFs of tuples fold at compile time, arrays of marginals
+# are summed dynamically (also on GPU arrays):
+@inline _sum_dofs(f, mar) = sum(f, mar)
+@inline _sum_dofs(f, mar::AbstractArray) = mapreduce(_dynamic_dof ∘ f, +, mar; init = 0)
+@inline _dynamic_dof(n::IntegerLike) = dynamic(n)
+@inline _dynamic_dof(nodof::AbstractNoDOF) = nodof
 
 function checked_arg(μ::ProductMeasure{<:NTuple{N,Any}}, x::NTuple{N,Any}) where {N}
     map(checked_arg, marginals(μ), x)
@@ -445,3 +452,28 @@ function _marginals_from_std_with_rest(::Type{S}, νs::AbstractArray{M}, z::Abst
         return [y for y in ys_any], z_rest
     end
 end
+
+# Batched transport of array products with scalar-variate marginals in one
+# broadcast, the marginals align with the leading dimension of the batch:
+
+function batched_transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, X::AbstractArray) where {S<:StdMeasure,M}
+    _array_product_batched_to_std(S, μ, X, mspace_flatsize(M), Val(isconcretetype(M)))
+end
+function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Tuple{}, ::Val{true}) where {S}
+    _check_flatsize(X, maybestatic_size(marginals(μ)))
+    _as_stream_batch(broadcast(_ToStd{S}(), marginals(μ), X), maybestatic_size(marginals(μ)))
+end
+function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Any, ::Val) where {S}
+    _batched_to_std(S, μ, X, mspace_flatsize(μ))
+end
+
+function batched_transport_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, Z::AbstractArray) where {S<:StdMeasure,M}
+    _array_product_batched_from_std(S, μ, Z, mspace_flatsize(M), Val(isconcretetype(M)))
+end
+function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Tuple{}, ::Val{true}) where {S}
+    mar = marginals(μ)
+    broadcast(_FromStd{S}(), mar, reshape(Z, (map(dynamic, maybestatic_size(mar))..., Base.tail(size(Z))...)))
+end
+function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Any, ::Val) where {S}
+    _batched_from_std(S, μ, Z, mspace_flatsize(μ))
+end
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 85b6ad0f..2e6a42b2 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -211,6 +211,28 @@ _pushfwd_dof(::Type{MU}, ::Type{<:Tuple{Any,Real}}, dof) where {MU} = dof
     return ν.f(x), z_rest
 end
 
+# Batched transport for pushforwards of measures with scalar variates, the
+# functions apply elementwise then:
+function batched_transport_to_std(::Type{S}, ν::PushforwardMeasure, Y::AbstractArray) where {S<:StdMeasure}
+    _pushfwd_batched_to_std(S, ν, Y, mspace_flatsize(ν.origin))
+end
+@inline function _pushfwd_batched_to_std(::Type{S}, ν::PushforwardMeasure, Y::AbstractArray, ::Tuple{}) where {S}
+    batched_transport_to_std(S, ν.origin, broadcast(ν.finv, Y))
+end
+@inline function _pushfwd_batched_to_std(::Type{S}, ν::PushforwardMeasure, Y::AbstractArray, ::Any) where {S}
+    _batched_to_std(S, ν, Y, mspace_flatsize(ν))
+end
+
+function batched_transport_from_std(::Type{S}, ν::PushforwardMeasure, Z::AbstractArray) where {S<:StdMeasure}
+    _pushfwd_batched_from_std(S, ν, Z, mspace_flatsize(ν.origin))
+end
+@inline function _pushfwd_batched_from_std(::Type{S}, ν::PushforwardMeasure, Z::AbstractArray, ::Tuple{}) where {S}
+    broadcast(ν.f, batched_transport_from_std(S, ν.origin, Z))
+end
+@inline function _pushfwd_batched_from_std(::Type{S}, ν::PushforwardMeasure, Z::AbstractArray, ::Any) where {S}
+    _batched_from_std(S, ν, Z, mspace_flatsize(ν))
+end
+
 massof(m::PushforwardMeasure) = massof(m.origin)
 
 function Base.rand(rng::AbstractRNG, ::Type{T}, ν::PushforwardMeasure) where {T}
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index 27f3504e..b964434a 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -79,3 +79,12 @@ insupport(μ::WeightedMeasure, x) = insupport(μ.base, x)
     transport_to_std_with_rest(S, basemeasure(μ), x)
 @inline transport_from_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, z::AbstractVector) where {S<:StdMeasure} =
     transport_from_std_with_rest(S, basemeasure(μ), z)
+
+@inline batched_transport_to_std(::Type{S}, μ::AbstractWeightedMeasure, X::AbstractArray) where {S<:StdMeasure} =
+    batched_transport_to_std(S, basemeasure(μ), X)
+@inline batched_transport_from_std(::Type{S}, μ::AbstractWeightedMeasure, Z::AbstractArray) where {S<:StdMeasure} =
+    batched_transport_from_std(S, basemeasure(μ), Z)
+@inline batched_transport_to_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, X::AbstractArray) where {S<:StdMeasure} =
+    batched_transport_to_std_with_rest(S, basemeasure(μ), X)
+@inline batched_transport_from_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, Z::AbstractArray) where {S<:StdMeasure} =
+    batched_transport_from_std_with_rest(S, basemeasure(μ), Z)
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 5b593da4..c5da71d4 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -56,3 +56,14 @@ end
 @inline transport_to_std(::Type{S}, ::Dirac, x) where {S<:StdMeasure} = SVector{0,Bool}()
 @inline transport_from_std(::Type{S}, μ::Dirac, z::AbstractVector) where {S<:StdMeasure} = μ.x
 @inline transport_from_std_with_rest(::Type{S}, μ::Dirac, z::AbstractVector) where {S<:StdMeasure} = μ.x, z
+
+function batched_transport_to_std(::Type{S}, μ::Dirac, X::AbstractArray) where {S<:StdMeasure}
+    n = length(_value_flatsize(μ.x))
+    similar(X, Bool, (0, ntuple(i -> size(X, n + i), Val(ndims(X) - n))...))
+end
+
+function batched_transport_from_std(::Type{S}, μ::Dirac, Z::AbstractArray) where {S<:StdMeasure}
+    X = similar(Z, eltype(μ.x), (size(μ.x)..., Base.tail(size(Z))...))
+    X .= μ.x
+    return X
+end
diff --git a/src/standard/stdconvert.jl b/src/standard/stdconvert.jl
index 6a5353ac..06fdd9c7 100644
--- a/src/standard/stdconvert.jl
+++ b/src/standard/stdconvert.jl
@@ -55,3 +55,9 @@ const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}}
 function transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, x) where {NU<:StdMeasure,MU<:StdMeasure}
     _pwr_variate(ν, maybestatic_reshape(stdconvert(NU, MU, x), mspace_flatsize(ν)))
 end
+
+function batched_transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, X::AbstractArray) where {NU<:StdMeasure,MU<:StdMeasure}
+    n_μ = length(mspace_flatsize(μ))
+    batch_dims = ntuple(i -> size(X, n_μ + i), Val(ndims(X) - n_μ))
+    reshape(stdconvert(NU, MU, X), (map(dynamic, _size_dims(mspace_flatsize(ν)))..., batch_dims...))
+end
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index 513aad8f..316d0d40 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -89,3 +89,6 @@ end
 end
 
 @inline _stdmeasure_rank(::Type{<:StdMeasure}) = 0
+
+@inline batched_transport_to_std(::Type{S}, ::S, X::AbstractArray) where {S<:StdMeasure} = _as_stdstream_batch(X)
+@inline batched_transport_from_std(::Type{S}, ::S, Z::AbstractArray) where {S<:StdMeasure} = _drop_stdstream_dim(Z)
diff --git a/src/transport-batched.jl b/src/transport-batched.jl
new file mode 100644
index 00000000..d14ea7ca
--- /dev/null
+++ b/src/transport-batched.jl
@@ -0,0 +1,156 @@
+# Batched transport over flat batches of variates: the leading dimensions of
+# a batch are the variate dimensions (see `mspace_flatsize`), all further
+# dimensions are batch dimensions. Streams of standard variates are batches
+# `(dof, batch dims...)`, consumed along their first dimension.
+
+"""
+    MeasureBase.batched_transport_to_std(::Type{S}, μ, X::AbstractArray)
+
+Batched form of [`MeasureBase.transport_to_std`](@ref): transports the
+flat batch `X` of variates of `μ` to a batch `(getdof(μ), batch dims...)`
+of variates of the standard measure type `S`.
+
+The default implementation broadcasts the point transport for measures
+with scalar variates and maps it over the variate slices of `X`
+otherwise.
+"""
+function batched_transport_to_std end
+
+function batched_transport_to_std(::Type{S}, μ, X::AbstractArray) where {S<:StdMeasure}
+    _batched_to_std(S, μ, X, mspace_flatsize(μ))
+end
+
+@inline function _batched_to_std(::Type{S}, μ, X::AbstractArray, ::Tuple{}) where {S}
+    _as_stdstream_batch(broadcast(Base.Fix1(_ToStd{S}(), μ), X))
+end
+
+function _batched_to_std(::Type{S}, μ, X::AbstractArray, sz::SizeLike) where {S}
+    _check_flatsize(X, sz)
+    stacked(map(Base.Fix1(_ToStd{S}(), μ), sliced(X, Val(length(sz)))))
+end
+
+function _batched_to_std(::Type{S}, μ, ::AbstractArray, ::NoMSpaceElementSize) where {S}
+    throw(ArgumentError("Batched transport requires measures of type $(nameof(typeof(μ))) to have a known variate size"))
+end
+
+# Standard variates of scalar-variate measures form the first dimension:
+@inline _as_stdstream_batch(Z::AbstractArray) = reshape(Z, (1, size(Z)...))
+@inline _drop_stdstream_dim(Z::AbstractArray) = reshape(Z, Base.tail(size(Z)))
+
+
+"""
+    MeasureBase.batched_transport_from_std(::Type{S}, μ, Z::AbstractArray)
+
+Batched form of [`MeasureBase.transport_from_std`](@ref): transports the
+batch `Z` of variates of the standard measure type `S`, of size
+`(getdof(μ), batch dims...)`, to a flat batch of variates of `μ`.
+"""
+function batched_transport_from_std end
+
+function batched_transport_from_std(::Type{S}, μ, Z::AbstractArray) where {S<:StdMeasure}
+    _batched_from_std(S, μ, Z, mspace_flatsize(μ))
+end
+
+@inline function _batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Tuple{}) where {S}
+    broadcast(Base.Fix1(_FromStd{S}(), μ), _drop_stdstream_dim(Z))
+end
+
+function _batched_from_std(::Type{S}, μ, Z::AbstractArray, sz::SizeLike) where {S}
+    xs = map(Base.Fix1(_FromStd{S}(), μ), sliced(Z, Val(1)))
+    reshape(stacked(xs), (map(dynamic, _size_dims(sz))..., Base.tail(size(Z))...))
+end
+
+function _batched_from_std(::Type{S}, μ, ::AbstractArray, ::NoMSpaceElementSize) where {S}
+    throw(ArgumentError("Batched transport requires measures of type $(nameof(typeof(μ))) to have a known variate size"))
+end
+
+
+"""
+    MeasureBase.batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray)
+
+Batched form of [`MeasureBase.transport_to_std_with_rest`](@ref) for a
+batch `X` of flat vector streams (first dimension along the streams).
+
+Returns a tuple `(Z, X_μ, X_rest)` of the batch of standard variates, the
+rows consumed from the streams and the unconsumed rest of the streams.
+"""
+function batched_transport_to_std_with_rest end
+
+function batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray) where {S<:StdMeasure}
+    X_μ, X_rest = _batched_consume(X, mspace_flatsize(μ))
+    return batched_transport_to_std(S, μ, X_μ), X_μ, X_rest
+end
+
+
+"""
+    MeasureBase.batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray)
+
+Batched form of [`MeasureBase.transport_from_std_with_rest`](@ref) for a
+batch `Z` of streams of standard variates (first dimension along the
+streams).
+
+Returns a tuple `(X, Z_rest)` of the flat batch of variates of `μ` and the
+unconsumed rest of the streams.
+"""
+function batched_transport_from_std_with_rest end
+
+function batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray) where {S<:StdMeasure}
+    _batched_from_std_with_rest_bydof(S, μ, Z, fast_dof(μ))
+end
+
+function _batched_from_std_with_rest_bydof(::Type{S}, μ, Z::AbstractArray, n::IntegerLike) where {S}
+    Z_μ, Z_rest = _batched_split(Z, dynamic(n))
+    return batched_transport_from_std(S, μ, Z_μ), Z_rest
+end
+
+function _batched_from_std_with_rest_bydof(::Type{S}, μ, ::AbstractArray, ::AbstractNoDOF) where {S}
+    throw(ArgumentError("Batched transport from standard measures requires measures of type $(nameof(typeof(μ))) to implement MeasureBase.batched_transport_from_std_with_rest"))
+end
+
+# A flat batch of variates as a batch of streams, the variate dimensions
+# merged into the first dimension:
+@inline function _as_stream_batch(X::AbstractArray, sz::SizeLike)
+    n = length(sz)
+    batch_dims = ntuple(i -> size(X, n + i), Val(ndims(X) - n))
+    reshape(X, (dynamic(size2length(sz)), batch_dims...))
+end
+
+
+"""
+    MeasureBase.batched_transport_def(ν, μ, X::AbstractArray)
+
+Transport the flat batch `X` of variates of `μ` to a flat batch of
+variates of `ν`, via the standard measure type the preferences of `ν` and
+`μ` promote to. Specialize for pairs of measure types with a direct
+batched transport.
+"""
+function batched_transport_def end
+
+function batched_transport_def(ν, μ, X::AbstractArray)
+    S = _transport_pivot(ν, μ)
+    Z = batched_transport_to_std(S, μ, X)
+    Y, Z_rest = batched_transport_from_std_with_rest(S, ν, Z)
+    if size(Z_rest, 1) != 0
+        throw(ArgumentError("Degrees of freedom of source and target measure of a transport don't match"))
+    end
+    return Y
+end
+
+
+# Broadcasting a transport function over an array of variates with flat
+# storage, or over the flat storage of a batch, transports the batch as a
+# whole. Variates of the target measure come out in their flat form.
+function Broadcast.broadcasted(f::TransportFunction, X::AbstractArray)
+    _broadcast_transport(f, X, _flat_storage(X), mspace_flatsize(f.μ), mspace_flatsize(f.ν))
+end
+
+function _broadcast_transport(f::TransportFunction, X, X_flat::AbstractArray, sz_μ::SizeLike, sz_ν::SizeLike)
+    _check_flatsize(X_flat, sz_μ)
+    Y_flat = batched_transport_def(f.ν, f.μ, X_flat)
+    return _batch_variates(Y_flat, sz_ν)
+end
+
+_broadcast_transport(f::TransportFunction, X, ::Any, ::Any, ::Any) = map(f, X)
+
+@inline _batch_variates(Y::AbstractArray, ::Tuple{}) = Y
+@inline _batch_variates(Y::AbstractArray, sz::SizeLike) = sliced(Y, Val(length(sz)))
diff --git a/test/distributions/test_transport.jl b/test/distributions/test_transport.jl
index f7515e18..c8d719b2 100644
--- a/test/distributions/test_transport.jl
+++ b/test/distributions/test_transport.jl
@@ -5,6 +5,7 @@ using Test
 using LinearAlgebra
 using InverseFunctions, ChangesOfVariables
 using Distributions, ArraysOfArrays
+using ArraysOfArrays: sliced, flatview
 using StableRNGs
 using LogExpFunctions: logit
 import ForwardDiff, Zygote
@@ -219,6 +220,22 @@ include("getjacobian.jl")
         @test transport_to(m, m2)(y) ≈ x
     end
 
+    @testset "batched transport" begin
+        mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3])
+        f = transport_to(StdNormal()^2, mvn)
+        X = rand(StableRNG(789990641), mvn, 6)
+        Y = f.(sliced(X, Val(1)))
+        @test flatview(Y) ≈ stack(map(f, eachcol(X)))
+        @test flatview(inverse(f).(Y)) ≈ X
+        g = transport_to(StdNormal(), Weibull(0.7, 1.3))
+        x = rand(StableRNG(789990641), Weibull(0.7, 1.3), 10)
+        @test g.(x) ≈ map(g, x)
+        pd = product_distribution([Weibull(0.7), Exponential(1.3), Normal(0.5, 2.0)])
+        h = transport_to(StdNormal()^3, asmeasure(pd))
+        Xp = rand(StableRNG(789990641), pd, 5)
+        @test stack(h.(sliced(Xp, Val(1)))) ≈ stack(map(h, eachcol(Xp)))
+    end
+
     @testset "MvNormal covariance representations" begin
         for Σ in [PDMats.ScalMat(3, 2.5), PDMats.PDiagMat([0.5, 1.0, 2.5]), Diagonal([0.5, 1.0, 2.5])]
             mvn = MvNormal([0.2, -0.4, 0.6], Σ)
diff --git a/test/reactant/runtests.jl b/test/reactant/runtests.jl
index 182034b8..1357cf8d 100644
--- a/test/reactant/runtests.jl
+++ b/test/reactant/runtests.jl
@@ -11,7 +11,7 @@ using MeasureBase
 using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic, Lebesgue, Dirac
 using MeasureBase: logdensities, logdensity_rel, weightedmeasure, superpose, restrict, mintegrate_exp
 using MeasureBase: mcombine
-using ArraysOfArrays: VectorOfSimilarVectors, sliced
+using ArraysOfArrays: VectorOfSimilarVectors, sliced, flatview
 using Distributions: Normal, Exponential, Uniform, Beta
 
 Reactant.set_default_backend("cpu")
@@ -88,6 +88,13 @@ _plain(x::Number) = Float64(x)
         test_traced(z -> transport_to(mc, StdNormal()^3)(z), randn(3))
     end
 
+    @testset "batched transport" begin
+        test_traced(X -> transport_to(StdNormal(), StdUniform()).(X), rand(10))
+        test_traced(X -> flatview(transport_to(StdExponential()^3, StdNormal()^3).(sliced(X, Val(1)))), X)
+        mc = mcombine(vcat, StdNormal()^2, StdUniform()^1)
+        test_traced(X -> flatview(transport_to(StdLogistic()^3, mc).(sliced(X, Val(1)))), vcat(X[1:2, :], rand(1, 20)))
+    end
+
     @testset "transport" begin
         test_traced(x -> transport_to(StdUniform(), StdNormal()).(x), x)
         test_traced(x -> transport_to(StdNormal(), StdUniform()).(x), rand(10))
diff --git a/test/transport.jl b/test/transport.jl
index 5a68d6b1..b239de90 100644
--- a/test/transport.jl
+++ b/test/transport.jl
@@ -5,9 +5,12 @@ using MeasureBase: StdUniform, StdExponential, StdLogistic, StdNormal
 using MeasureBase: Dirac, Half, restrict, mbind, productmeasure, pushfwd
 using MeasureBase: transport_to_std, transport_from_std, transport_from_std_with_rest
 using InverseFunctions: inverse
+using MeasureBase: weightedmeasure, mcombine
 using StaticArrays: SVector
 using Static: static
 using LogExpFunctions: logit
+using ArraysOfArrays: sliced, flatview
+using JLArrays
 
 @testset "transport_to" begin
     for (f, μ) in [
@@ -133,6 +136,64 @@ using LogExpFunctions: logit
         @test_throws ArgumentError transport_to(μ, StdUniform())(0.5)
     end
 
+    @testset "batched transport" begin
+        f = transport_to(StdNormal(), StdUniform())
+        X = rand(7)
+        @test f.(X) ≈ map(f, X)
+        @test inverse(f).(f.(X)) ≈ X
+        @test eltype(f.(rand(Float32, 5))) == Float32
+
+        g = transport_to(StdExponential()^3, StdNormal()^3)
+        Xn = randn(3, 5)
+        Yn = g.(sliced(Xn, Val(1)))
+        @test Yn isa AbstractVector && length(Yn) == 5
+        @test flatview(Yn) ≈ stack(map(g, eachcol(Xn)))
+        @test flatview(g.(Xn)) ≈ flatview(Yn)
+        @test flatview(inverse(g).(Yn)) ≈ Xn
+        Xv = [randn(3) for _ in 1:4]
+        @test g.(Xv) == map(g, Xv)
+
+        h = transport_to(StdUniform()^(2, 3), (StdNormal()^2)^3)
+        Xh = randn(2, 3, 4)
+        @test flatview(h.(Xh)) ≈ stack([h(Xh[:, :, i]) for i in 1:4])
+        @test flatview(inverse(h).(h.(Xh))) ≈ Xh
+
+        P = MeasureBase.ProductMeasure([weightedmeasure(log(i), StdNormal()) for i in 1:3])
+        p = transport_to(StdUniform()^3, P)
+        Xp = randn(3, 6)
+        Yp = p.(sliced(Xp, Val(1)))
+        @test flatview(Yp) ≈ stack(map(p, eachcol(Xp)))
+        @test flatview(inverse(p).(Yp)) ≈ Xp
+
+        mc = mcombine(vcat, StdNormal()^2, StdUniform()^3)
+        c = transport_to(StdExponential()^5, mc)
+        Xc = vcat(randn(2, 4), rand(3, 4))
+        Yc = c.(sliced(Xc, Val(1)))
+        @test flatview(Yc) ≈ stack(map(c, eachcol(Xc)))
+        @test flatview(inverse(c).(Yc)) ≈ Xc
+        cd = transport_to(mcombine(vcat, Dirac(0.5), StdUniform()^2), StdNormal()^2)
+        @test flatview(cd.(randn(2, 3)))[1, :] == fill(0.5, 3)
+
+        pf = transport_to(StdUniform(), pushfwd(exp, StdNormal()))
+        Xe = exp.(randn(8))
+        @test pf.(Xe) ≈ map(pf, Xe)
+        @test inverse(pf).(pf.(Xe)) ≈ Xe
+
+        w = transport_to(StdLogistic()^2, weightedmeasure(0.3, StdNormal()^2))
+        Xw = randn(2, 5)
+        @test flatview(w.(sliced(Xw, Val(1)))) ≈ stack(map(w, eachcol(Xw)))
+
+        JLArrays.allowscalar(false)
+        Xj = JLArray(Xn)
+        Yj = g.(sliced(Xj, Val(1)))
+        @test flatview(Yj) isa JLArray
+        @test Array(flatview(Yj)) ≈ flatview(Yn)
+        @test Array(flatview(c.(sliced(JLArray(Xc), Val(1))))) ≈ flatview(Yc)
+        Pj = MeasureBase.ProductMeasure(JLArray([weightedmeasure(log(i), StdNormal()) for i in 1:3]))
+        pj = transport_to(StdUniform()^3, Pj)
+        @test Array(flatview(pj.(sliced(JLArray(Xp), Val(1))))) ≈ flatview(Yp)
+    end
+
     @testset "transport for products" begin
         test_transport(
             StdUniform()^(2, 2),

From 476041403a75f939e1cd98e0676f8f813591ac36 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 01:53:21 +0200
Subject: [PATCH 089/122] Generate random variates through generative contexts
 in flat batches

Random variate generation now runs on a GenContext carrying the random
number generator, the numerical precision and the compute unit. Measure
types implement rand_impl for single variates and batched_rand_impl for
flat batches, with defaults that draw variates of the preferred standard
measure and transport them. Variates of powers are drawn as one flat
batch of the innermost base measure and reshaped into the layout of the
flat variate storage, so that they feed the batched density and
transport paths directly. Superpositions and spike mixtures select
between component batches branch-free. Standard measures have unit mass.

Created by generative AI.
---
 .../MeasureBaseDistributionsExt.jl            |   2 +-
 .../distribution_measure.jl                   |  14 +-
 src/MeasureBase.jl                            |   5 +-
 src/combinators/bind.jl                       |  14 +-
 src/combinators/combined.jl                   |  21 +--
 src/combinators/half.jl                       |   5 +-
 src/combinators/power.jl                      |  48 ++++--
 src/combinators/product.jl                    |  53 +------
 src/combinators/spikemixture.jl               |   9 +-
 src/combinators/superpose.jl                  |  34 ++++-
 src/combinators/transformedmeasure.jl         |  17 ++-
 src/combinators/weighted.jl                   |   6 +-
 src/mass-interface.jl                         |   2 +
 src/primitives/dirac.jl                       |   3 +-
 src/rand.jl                                   | 141 +++++++++++++++---
 src/standard/stdexponential.jl                |   3 +-
 src/standard/stdlogistic.jl                   |   5 +-
 src/standard/stdmeasure.jl                    |   2 +
 src/standard/stdnormal.jl                     |   3 +-
 src/standard/stduniform.jl                    |   3 +-
 test/rand.jl                                  | 109 ++++++++++++++
 test/runtests.jl                              |   2 +
 22 files changed, 365 insertions(+), 136 deletions(-)
 create mode 100644 test/rand.jl

diff --git a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
index 586ae186..83a20c65 100644
--- a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
+++ b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
@@ -41,7 +41,7 @@ import PDMats
 using IrrationalConstants: log2π, invsqrt2π
 using LogExpFunctions: logistic
 
-using HeterogeneousComputing: real_numtype
+using HeterogeneousComputing: real_numtype, GenContext, get_rng, get_precision
 
 using Static: True, False, StaticInt, static, dynamic
 using StaticThings: asnonstatic
diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index 80c6f53e..652826ad 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -15,7 +15,11 @@ const DistributionMeasure{F<:VariateForm,S<:ValueSupport,D<:Distribution{F,S}} =
 @inline Base.convert(::Type{Distribution{F,S}}, m::DistributionMeasure{F,S}) where {F<:VariateForm,S<:ValueSupport} = Distribution(m)
 
 
-Base.rand(rng::AbstractRNG, ::Type{T}, m::DistributionMeasure) where {T<:Real} = convert_realtype(T, rand(m.obj))
+MeasureBase.rand_impl(ctx::GenContext, m::DistributionMeasure) =
+    convert_realtype(get_precision(ctx), rand(get_rng(ctx), m.obj))
+
+MeasureBase.batched_rand_impl(ctx::GenContext, m::DistributionMeasure, sz::Dims) =
+    _flat_powrand(get_rng(ctx), get_precision(ctx), m.obj, sz)
 
 function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution{<:ArrayLikeVariate{0}}, sz::Dims) where {T<:Real}
     convert_realtype(T, reshape(rand(rng, d, prod(sz)), sz...))
@@ -33,14 +37,6 @@ function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution, sz::Dims) w
     flatview(ArrayOfSimilarArrays(convert_realtype(T, rand(rng, d, sz))))
 end
 
-function Base.rand(rng::AbstractRNG, ::Type{T}, m::PowerMeasure{<:DistributionMeasure{<:ArrayLikeVariate{0}}, NTuple{N,Base.OneTo{Int}}}) where {T<:Real,N}
-    _flat_powrand(rng, T, m.parent.obj, map(length, m.axes))
-end
-
-function Base.rand(rng::AbstractRNG, ::Type{T}, m::PowerMeasure{<:DistributionMeasure{<:ArrayLikeVariate{M}}, NTuple{N,Base.OneTo{Int}}}) where {T<:Real,M,N}
-    flat_data = _flat_powrand(rng, T, m.parent.obj, map(length, m.axes))
-    ArrayOfSimilarArrays{T,M,N}(flat_data)
-end
 
 
 @inline DensityInterface.densityof(m::DistributionMeasure) = densityof(m.obj)
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 7e62d783..5aa2eb09 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -59,6 +59,8 @@ using StaticThings:
 
 import HeterogeneousComputing
 using HeterogeneousComputing: real_numtype
+using HeterogeneousComputing:
+    GenContext, AbstractComputeUnit, CPUnit, get_rng, get_precision, get_compute_unit, allocate_array
 
 using ArraysOfArrays:
     ArrayOfSimilarArrays, VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays,
@@ -186,6 +188,7 @@ include("mspace.jl")
 include("getdof.jl")
 include("standard/stdmeasure.jl")
 include("transport.jl")
+include("rand.jl")
 include("proxies.jl")
 include("parameterized.jl")
 include("domains.jl")
@@ -228,8 +231,6 @@ include("combinators/half.jl")
 
 #include("implicitmaps.jl")
 
-include("rand.jl")
-
 include("measure_operators.jl")
 
 include("interface.jl")
diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl
index 7a2bbed8..2c24ca7f 100644
--- a/src/combinators/bind.jl
+++ b/src/combinators/bind.jl
@@ -334,17 +334,15 @@ function logdensityof_with_rest(μ::_BindBy{typeof(merge)}, x::NamedTuple)
 end
 
 
-function Base.rand(rng::Random.AbstractRNG, ::Type{T}, μ::Bind) where {T<:Real}
-    a = rand(rng, T, μ.α)
-    b = rand(rng, T, _get_β_a(μ, a))
+function rand_impl(ctx::GenContext, μ::Bind)
+    a = rand_impl(ctx, μ.α)
+    b = rand_impl(ctx, _get_β_a(μ, a))
     return μ.f_c(a, b)
 end
 
-function Base.rand(rng::Random.AbstractRNG, μ::Bind)
-    a = rand(rng, μ.α)
-    b = rand(rng, _get_β_a(μ, a))
-    return μ.f_c(a, b)
-end
+# The secondary measure depends on the primary variate, so batches are
+# generated variate by variate:
+batched_rand_impl(ctx::GenContext, μ::Bind, sz::Dims) = _batched_rand_pointwise(ctx, μ, sz)
 
 
 # Transport consumes the variate parts of the primary and secondary
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index 7d9ecedb..68cd16d6 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -216,16 +216,19 @@ function logdensityof_with_rest(μ::CombinedMeasure{typeof(merge)}, x::NamedTupl
 end
 
 
-function Base.rand(rng::Random.AbstractRNG, ::Type{T}, μ::CombinedMeasure) where {T<:Real}
-    a = rand(rng, T, μ.α)
-    b = rand(rng, T, μ.β)
-    return μ.f_c(a, b)
-end
+rand_impl(ctx::GenContext, μ::CombinedMeasure) = μ.f_c(rand_impl(ctx, μ.α), rand_impl(ctx, μ.β))
 
-function Base.rand(rng::Random.AbstractRNG, μ::CombinedMeasure)
-    a = rand(rng, μ.α)
-    b = rand(rng, μ.β)
-    return μ.f_c(a, b)
+# Batches of vcat-combined measures are concatenated along the streams:
+function batched_rand_impl(ctx::GenContext, μ::CombinedMeasure{typeof(vcat)}, sz::Dims)
+    _combined_batched_rand(ctx, μ, sz, mspace_flatsize(μ.α), mspace_flatsize(μ.β))
+end
+function _combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::Dims, sz_a::SizeLike, sz_b::SizeLike)
+    A = _as_stream_batch(batched_rand_impl(ctx, μ.α, sz), sz_a)
+    B = _as_stream_batch(batched_rand_impl(ctx, μ.β, sz), sz_b)
+    return vcat(A, B)
+end
+function _combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::Dims, ::Any, ::Any)
+    _batched_rand_pointwise(ctx, μ, sz)
 end
 
 
diff --git a/src/combinators/half.jl b/src/combinators/half.jl
index 899bf52f..3873a86c 100644
--- a/src/combinators/half.jl
+++ b/src/combinators/half.jl
@@ -20,9 +20,8 @@ unhalf(μ::Half) = μ.parent
     weightedmeasure(logtwo, basemeasure(unhalf(μ)))
 end
 
-function Base.rand(rng::AbstractRNG, ::Type{T}, μ::Half) where {T}
-    return abs(rand(rng, T, unhalf(μ)))
-end
+@inline rand_impl(ctx::GenContext, μ::Half) = abs(rand_impl(ctx, unhalf(μ)))
+@inline batched_rand_impl(ctx::GenContext, μ::Half, sz::Dims) = abs.(batched_rand_impl(ctx, unhalf(μ), sz))
 
 function logdensityof_impl(μ::Half, x)
     ld = logdensityof(unhalf(μ), x) - loghalf
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 76afc6a5..026a65c5 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -51,28 +51,44 @@ function Pretty.tile(μ::PowerMeasure)
     return Pretty.pair_layout(arg1, arg2; sep = " ^ ")
 end
 
-# ToDo: Make rand return static arrays for statically-sized power measures.
-
 function _cartidxs(axs::Tuple{Vararg{AbstractUnitRange,N}}) where {N}
     CartesianIndices(map(asnonstatic, axs))
 end
 
-function Base.rand(
-    rng::AbstractRNG,
-    ::Type{T},
-    d::PowerMeasure{M},
-) where {T,M<:AbstractMeasure}
-    axs, base_d = pwr_axes(d), pwr_base(d)
-    map(_cartidxs(axs)) do _
-        rand(rng, T, base_d)
-    end
+# Variates of powers are generated as one flat batch of variates of the
+# innermost base measure, in the layout of the flat variate storage:
+
+rand_impl(ctx::GenContext, μ::PowerMeasure) = _pwr_rand(ctx, μ, mspace_flatsize(μ))
+
+function _pwr_rand(ctx::GenContext, μ::PowerMeasure, sz_flat::SizeLike)
+    ν, _ = _pwr_unwrap(μ)
+    _pwr_variate(μ, batched_rand_impl(ctx, ν, _pwr_batch_dims(sz_flat, mspace_flatsize(ν))))
 end
 
-function Base.rand(rng::AbstractRNG, ::Type{T}, d::PowerMeasure) where {T}
-    axs, base_d = pwr_axes(d), pwr_base(d)
-    map(_cartidxs(axs)) do _
-        rand(rng, base_d)
-    end
+function _pwr_rand(ctx::GenContext, μ::PowerMeasure, ::NoMSpaceElementSize)
+    ν = pwr_base(μ)
+    map(_ -> rand_impl(ctx, ν), _cartidxs(pwr_axes(μ)))
+end
+
+# The power dimensions of a flat size, after the flat dimensions of the
+# innermost base measure:
+@inline function _pwr_batch_dims(sz_flat::SizeLike, sz_base::SizeLike)
+    dims = map(dynamic, _size_dims(sz_flat))
+    n = length(sz_base)
+    ntuple(i -> dims[n + i], Val(length(dims) - n))
+end
+
+function batched_rand_impl(ctx::GenContext, μ::PowerMeasure, sz::Dims)
+    _pwr_batched_rand(ctx, μ, sz, mspace_flatsize(μ))
+end
+
+function _pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::Dims, sz_flat::SizeLike)
+    ν, _ = _pwr_unwrap(μ)
+    batched_rand_impl(ctx, ν, (_pwr_batch_dims(sz_flat, mspace_flatsize(ν))..., sz...))
+end
+
+function _pwr_batched_rand(::GenContext, μ::PowerMeasure, ::Dims, ::NoMSpaceElementSize)
+    throw(ArgumentError("Batched random variate generation for powers of measures of type $(nameof(typeof(pwr_base(μ)))) requires a known variate size"))
 end
 
 marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes)
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index b63821ca..fc5c37fa 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -28,49 +28,10 @@ Base.size(μ::AbstractProductMeasure) = size(marginals(μ))
 
 basemeasure(d::AbstractProductMeasure) = productmeasure(map(basemeasure, marginals(d)))
 
-function Base.rand(rng::AbstractRNG, ::Type{T}, d::AbstractProductMeasure) where {T}
-    mar = marginals(d)
-    _rand_product(rng, T, mar, eltype(mar))
-end
-
-function _rand_product(
-    rng::AbstractRNG,
-    ::Type{T},
-    mar,
-    ::Type{M},
-) where {T,M<:AbstractMeasure}
-    map(mar) do dⱼ
-        rand(rng, T, dⱼ)
-    end
-end
-
-function _rand_product(
-    rng::AbstractRNG,
-    ::Type{T},
-    mar::ReadonlyMappedArray,
-    ::Type{M},
-) where {T,M<:AbstractMeasure}
-    mappedarray(mar.data) do dⱼ
-        rand(rng, T, mar.f(dⱼ))
-    end |> collect
-end
-
-function _rand_product(rng::AbstractRNG, ::Type{T}, mar, ::Type{M}) where {T,M}
-    map(mar) do dⱼ
-        rand(rng, dⱼ)
-    end
-end
+rand_impl(ctx::GenContext, d::AbstractProductMeasure) = map(Base.Fix1(_marginal_rand, ctx), marginals(d))
 
-function _rand_product(
-    rng::AbstractRNG,
-    ::Type{T},
-    mar::ReadonlyMappedArray,
-    ::Type{M},
-) where {T,M}
-    mappedarray(mar.data) do dⱼ
-        rand(rng, mar.f(dⱼ))
-    end |> collect
-end
+@inline _marginal_rand(ctx::GenContext, m::AbstractMeasure) = rand_impl(ctx, m)
+@inline _marginal_rand(ctx::GenContext, d) = rand(get_rng(ctx), d)
 
 for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)]
     @eval @inline function $head(d::AbstractProductMeasure, x)
@@ -308,14 +269,6 @@ end
 export rand!
 using Random: rand!, GLOBAL_RNG
 
-function _rand(rng::AbstractRNG, ::Type{T}, d::ProductMeasure, mar::AbstractArray) where {T}
-    elT = typeof(rand(rng, T, first(mar)))
-
-    sz = size(mar)
-    x = Array{elT,length(sz)}(undef, sz)
-    rand!(rng, d, x)
-end
-
 @inline function insupport(d::AbstractProductMeasure, x::AbstractArray)
     _all_insupport(broadcast(_insupport_bool ∘ insupport, marginals(d), x))
 end
diff --git a/src/combinators/spikemixture.jl b/src/combinators/spikemixture.jl
index 64033840..440c67d5 100644
--- a/src/combinators/spikemixture.jl
+++ b/src/combinators/spikemixture.jl
@@ -33,8 +33,13 @@ function gentype(μ::SpikeMixture)
     gentype(μ.m)
 end
 
-function Base.rand(rng::AbstractRNG, T::Type, μ::SpikeMixture)
-    return (rand(rng, T) < μ.w) * rand(rng, T, μ.m)
+function rand_impl(ctx::GenContext, μ::SpikeMixture)
+    return (rand(get_rng(ctx), get_precision(ctx)) < μ.w) * rand_impl(ctx, μ.m)
+end
+
+function batched_rand_impl(ctx::GenContext, μ::SpikeMixture, sz::Dims)
+    X = batched_rand_impl(ctx, μ.m, sz)
+    return ifelse.(_rand_bulk(ctx, sz) .< μ.w, X, zero(eltype(X)))
 end
 
 testvalue(::Type{T}, μ::SpikeMixture) where {T} = zero(T)
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index cb8073fe..1c81886e 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -125,20 +125,40 @@ end
 
 basemeasure(μ::SuperpositionMeasure) = superpose(map(basemeasure, μ.components))
 
-function Base.rand(rng::AbstractRNG, ::Type{T}, μ::SuperpositionMeasure) where {T}
-    components = values(μ.components)
-    masses = map(massof, components)
+function _component_masses(μ::SuperpositionMeasure)
+    masses = map(massof, values(μ.components))
     total = sum(masses)
     total isa AbstractUnknownMass && throw(
         ArgumentError("Cannot sample from a superposition of measures of unknown mass"),
     )
-    threshold = rand(rng) * dynamic(total)
+    return map(dynamic, masses), dynamic(total)
+end
+
+function rand_impl(ctx::GenContext, μ::SuperpositionMeasure)
+    components = values(μ.components)
+    masses, total = _component_masses(μ)
+    threshold = rand(get_rng(ctx), get_precision(ctx)) * total
     csum = zero(threshold)
     for (mass, c) in zip(masses, components)
-        csum += dynamic(mass)
-        csum >= threshold && return rand(rng, T, c)
+        csum += mass
+        csum >= threshold && return rand_impl(ctx, c)
+    end
+    return rand_impl(ctx, last(components))
+end
+
+# Batches of superpositions draw a batch from each component and select
+# by mass, branch-free:
+function batched_rand_impl(ctx::GenContext, μ::SuperpositionMeasure, sz::Dims)
+    components = values(μ.components)
+    masses, total = _component_masses(μ)
+    thresholds = _rand_bulk(ctx, sz) .* total
+    X = batched_rand_impl(ctx, first(components), sz)
+    csum = first(masses)
+    for (mass, c) in Iterators.drop(zip(masses, components), 1)
+        X = ifelse.(thresholds .<= csum, X, batched_rand_impl(ctx, c, sz))
+        csum += mass
     end
-    return rand(rng, T, last(components))
+    return X
 end
 
 @inline function insupport(d::SuperpositionMeasure, x)
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 2e6a42b2..7db728e8 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -235,8 +235,21 @@ end
 
 massof(m::PushforwardMeasure) = massof(m.origin)
 
-function Base.rand(rng::AbstractRNG, ::Type{T}, ν::PushforwardMeasure) where {T}
-    return ν.f(rand(rng, T, ν.origin))
+rand_impl(ctx::GenContext, ν::PushforwardMeasure) = ν.f(rand_impl(ctx, ν.origin))
+
+# Batches of pushforwards apply the function to the variates of a batch of
+# the origin, elementwise for scalar variates:
+function batched_rand_impl(ctx::GenContext, ν::PushforwardMeasure, sz::Dims)
+    _pushfwd_batched_rand(ctx, ν, sz, mspace_flatsize(ν.origin))
+end
+@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, ::Tuple{})
+    broadcast(ν.f, batched_rand_impl(ctx, ν.origin, sz))
+end
+function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, sz_orig::SizeLike)
+    stacked(map(ν.f, sliced(batched_rand_impl(ctx, ν.origin, sz), Val(length(sz_orig)))))
+end
+@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, ::NoMSpaceElementSize)
+    _batched_rand_pointwise(ctx, ν, sz)
 end
 
 ###############################################################################
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index b964434a..b9a8a89b 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -30,9 +30,9 @@ end
     _lazy_add(_logweight_for(d.logweight, A), batched_logdensityof_impl(basemeasure(d), A))
 end
 
-function Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractWeightedMeasure) where {T}
-    rand(rng, T, basemeasure(μ))
-end
+@inline rand_impl(ctx::GenContext, μ::AbstractWeightedMeasure) = rand_impl(ctx, basemeasure(μ))
+@inline batched_rand_impl(ctx::GenContext, μ::AbstractWeightedMeasure, sz::Dims) =
+    batched_rand_impl(ctx, basemeasure(μ), sz)
 
 testvalue(::Type{T}, μ::AbstractWeightedMeasure) where {T} = testvalue(T, basemeasure(μ))
 
diff --git a/src/mass-interface.jl b/src/mass-interface.jl
index 33937498..548a941c 100644
--- a/src/mass-interface.jl
+++ b/src/mass-interface.jl
@@ -22,6 +22,8 @@ for T in (:UnknownFiniteMass, :UnknownMass)
     @eval begin
         Base.:+(::$T, ::$T) = $T()
         Base.:*(::$T, ::$T) = $T()
+        Base.:*(x::Real, ::$T) = isfinite(x) ? $T() : UnknownMass()
+        Base.:*(::$T, x::Real) = isfinite(x) ? $T() : UnknownMass()
         Base.:^(::$T, k::Real) = isfinite(k) ? $T() : UnknownMass()
         # Disambiguation:
         Base.:^(::$T, k::Integer) = isfinite(k) ? $T() : UnknownMass()
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index c5da71d4..1e542339 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -33,7 +33,8 @@ logdensityof_impl(μ::Dirac, x) = _checksupport(insupport(μ, x), zero(_logd_num
 logdensity_def(::Dirac, x::Number) = zero(float(typeof(x)))
 logdensity_def(::Dirac, x) = zero(_logd_numtype(x))
 
-Base.rand(::Random.AbstractRNG, T::Type, μ::Dirac) = μ.x
+@inline rand_impl(::GenContext, μ::Dirac) = μ.x
+@inline batched_rand_impl(ctx::GenContext, μ::Dirac, sz::Dims) = _const_batch(ctx, μ.x, sz)
 
 export dirac
 
diff --git a/src/rand.jl b/src/rand.jl
index f92cb16a..cf8a8a5e 100644
--- a/src/rand.jl
+++ b/src/rand.jl
@@ -1,24 +1,131 @@
-import Base
+# Random variate generation is parameterized by a `GenContext` carrying the
+# random number generator, the numerical precision and the compute unit.
+# Batches of variates are generated in flat form on the compute unit, single
+# variates of powers are drawn as one batch and reshaped.
 
-Base.rand(d::AbstractMeasure) = rand(Random.GLOBAL_RNG, Float64, d)
+"""
+    rand([rng::AbstractRNG], [T::Type{<:AbstractFloat}], μ::AbstractMeasure)
+    rand(ctx::GenContext, μ::AbstractMeasure)
 
-Base.rand(T::Type, μ::AbstractMeasure) = rand(Random.GLOBAL_RNG, T, μ)
+Generate a random variate of `μ`.
 
-Base.rand(rng::AbstractRNG, d::AbstractMeasure) = rand(rng, Float64, d)
+The generative context `ctx` (see `HeterogeneousComputing.GenContext`)
+determines the random number generator, the numerical precision (`Float64`
+by default) and the compute unit that array-valued variates are generated
+on. The variates of powers of measures are generated in one batch.
 
-@inline Random.rand!(d::AbstractMeasure, args...) = rand!(GLOBAL_RNG, d, args...)
+Measure types should specialize [`MeasureBase.rand_impl`](@ref) and
+[`MeasureBase.batched_rand_impl`](@ref) instead of `rand`.
+"""
+Base.rand(ctx::GenContext, μ::AbstractMeasure) = rand_impl(ctx, μ)
 
-# TODO: Make this work
-# function Base.rand(rng::AbstractRNG, ::Type{T}, d::AbstractMeasure) where {T}
-#     x = testvalue(d)
-#     rand!(d, x)
-# end
+Base.rand(μ::AbstractMeasure) = rand(GenContext{Float64}(), μ)
+Base.rand(rng::AbstractRNG, μ::AbstractMeasure) = rand(GenContext{Float64}(rng), μ)
+Base.rand(::Type{T}, μ::AbstractMeasure) where {T<:AbstractFloat} = rand(GenContext{T}(), μ)
+Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractMeasure) where {T<:AbstractFloat} = rand(GenContext{T}(rng), μ)
 
-# struct ArraySlot{A,I}
-#     arr::A
-#     i::I
-# end
+@inline Random.rand!(d::AbstractMeasure, args...) = rand!(Random.default_rng(), d, args...)
 
-# function rand!(rng::AbstractRNG, d::AbstractMeasure, x::ArraySlot)
-#     x.arr[x.i...] = rand(rng, d)
-# end
+
+"""
+    MeasureBase.rand_impl(ctx::GenContext, μ)
+
+Generate one random variate of `μ` in the generative context `ctx`.
+
+The default implementation draws a variate of the preferred standard
+measure of `μ` and transports it to `μ`. Measure types with a more direct
+way of generating variates specialize `rand_impl`, and should specialize
+[`MeasureBase.batched_rand_impl`](@ref) as well where batches can be
+generated in a more direct way, too.
+"""
+function rand_impl end
+
+function rand_impl(ctx::GenContext, μ)
+    _rand_via_std(ctx, μ, preferred_stdmeasure(μ), fast_dof(μ), mspace_flatsize(μ))
+end
+
+@inline function _rand_via_std(ctx::GenContext, μ, ::Type{S}, ::IntegerLike, ::Tuple{}) where {S<:StdMeasure}
+    transport_from_std(S, μ, rand_impl(ctx, S()))
+end
+@inline function _rand_via_std(ctx::GenContext, μ, ::Type{S}, n::IntegerLike, ::Any) where {S<:StdMeasure}
+    transport_from_std(S, μ, _rand_std(ctx, S, (dynamic(n),)))
+end
+@inline function _rand_via_std(ctx::GenContext, μ, ::Type{AnyStdMeasure}, n::IntegerLike, sz)
+    _rand_via_std(ctx, μ, StdUniform, n, sz)
+end
+function _rand_via_std(::GenContext, μ, ::Any, ::Any, ::Any)
+    throw(ArgumentError("Random variate generation is not implemented for measures of type $(nameof(typeof(μ))), define MeasureBase.rand_impl"))
+end
+
+
+"""
+    MeasureBase.batched_rand_impl(ctx::GenContext, μ, sz::Dims)
+
+Generate a batch of random variates of `μ` of batch size `sz` in flat
+form, an array of size `(flat variate dims..., sz...)` (see
+[`MeasureBase.mspace_flatsize`](@ref)).
+
+The default implementation draws a batch of variates of the preferred
+standard measure of `μ` and transports it to `μ`, or generates the
+variates one by one if `μ` has no standard transport.
+"""
+function batched_rand_impl end
+
+function batched_rand_impl(ctx::GenContext, μ, sz::Dims)
+    _batched_rand_via_std(ctx, μ, sz, preferred_stdmeasure(μ), fast_dof(μ), mspace_flatsize(μ))
+end
+
+function _batched_rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{S}, n::IntegerLike, ::SizeLike) where {S<:StdMeasure}
+    batched_transport_from_std(S, μ, _rand_std(ctx, S, (dynamic(n), sz...)))
+end
+@inline function _batched_rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{AnyStdMeasure}, n::IntegerLike, sz_flat::SizeLike)
+    _batched_rand_via_std(ctx, μ, sz, StdUniform, n, sz_flat)
+end
+@inline function _batched_rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{AnyStdMeasure}, n::IntegerLike, sz_flat::NoMSpaceElementSize)
+    _batched_rand_via_std(ctx, μ, sz, StdUniform, n, sz_flat)
+end
+function _batched_rand_via_std(ctx::GenContext, μ, sz::Dims, ::Any, ::Any, ::SizeLike)
+    _batched_rand_pointwise(ctx, μ, sz)
+end
+function _batched_rand_via_std(::GenContext, μ, ::Dims, ::Any, ::Any, ::NoMSpaceElementSize)
+    throw(ArgumentError("Batched random variate generation requires measures of type $(nameof(typeof(μ))) to have a known variate size"))
+end
+function _batched_rand_via_std(::GenContext, μ, ::Dims, ::Type{S}, ::IntegerLike, ::NoMSpaceElementSize) where {S<:StdMeasure}
+    throw(ArgumentError("Batched random variate generation requires measures of type $(nameof(typeof(μ))) to have a known variate size"))
+end
+
+function _batched_rand_pointwise(ctx::GenContext, μ, sz::Dims)
+    _stack_variates(map(_ -> rand_impl(ctx, μ), CartesianIndices(sz)))
+end
+
+@inline _stack_variates(xs::AbstractArray{<:Number}) = xs
+@inline _stack_variates(xs::AbstractArray{<:AbstractArray}) = stacked(xs)
+
+
+# Bulk draws of standard variates on the compute unit:
+
+@inline _rand_std(ctx::GenContext, ::Type{S}, dims::Dims) where {S<:StdMeasure} = batched_rand_impl(ctx, S(), dims)
+
+@inline _rand_bulk(ctx::GenContext, sz::Dims) = rand(ctx, sz)
+@inline _randn_bulk(ctx::GenContext, sz::Dims) = randn(ctx, sz)
+@inline _randexp_bulk(ctx::GenContext, sz::Dims) = _randexp_bulk(ctx, sz, get_compute_unit(ctx))
+@inline _randexp_bulk(ctx::GenContext, sz::Dims, ::CPUnit) = randexp(ctx, sz)
+# Not all compute units provide exponential draws, derive them from uniform draws then:
+@inline _randexp_bulk(ctx::GenContext, sz::Dims, ::AbstractComputeUnit) = -log1p.(-_rand_bulk(ctx, sz))
+
+# Test values use a constant RNG, which only draws single values:
+const _ConstantContext = GenContext{<:Any,<:Any,ConstantRNG}
+@inline _rand_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, rand(ConstantRNG(), get_precision(ctx)), sz)
+@inline _randn_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, randn(ConstantRNG(), get_precision(ctx)), sz)
+@inline _randexp_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, randexp(ConstantRNG(), get_precision(ctx)), sz)
+@inline _const_bulk(ctx::GenContext, x, sz::Dims) = fill!(allocate_array(ctx, typeof(x), sz), x)
+
+# A batch of copies of a constant variate:
+function _const_batch(ctx::GenContext, x, sz::Dims)
+    X = allocate_array(ctx, eltype(x), (size(x)..., sz...))
+    X .= x
+    return X
+end
+function _const_batch(ctx::GenContext, x::Number, sz::Dims)
+    fill!(allocate_array(ctx, typeof(x), sz), x)
+end
diff --git a/src/standard/stdexponential.jl b/src/standard/stdexponential.jl
index 1e10bb88..82abc200 100644
--- a/src/standard/stdexponential.jl
+++ b/src/standard/stdexponential.jl
@@ -15,4 +15,5 @@ end
 @inline transport_def(::StdUniform, μ::StdExponential, x) = -expm1(-x)
 @inline transport_def(::StdExponential, μ::StdUniform, x) = -log1p(-x)
 
-Base.rand(rng::Random.AbstractRNG, ::Type{T}, ::StdExponential) where {T} = randexp(rng, T)
+@inline rand_impl(ctx::GenContext, ::StdExponential) = randexp(get_rng(ctx), get_precision(ctx))
+@inline batched_rand_impl(ctx::GenContext, ::StdExponential, sz::Dims) = _randexp_bulk(ctx, sz)
diff --git a/src/standard/stdlogistic.jl b/src/standard/stdlogistic.jl
index b28dd618..fa6415a0 100644
--- a/src/standard/stdlogistic.jl
+++ b/src/standard/stdlogistic.jl
@@ -12,9 +12,8 @@ export StdLogistic
 @inline transport_def(::StdUniform, μ::StdLogistic, x) = logistic(x)
 @inline transport_def(::StdLogistic, μ::StdUniform, p) = logit(p)
 
-@inline function Base.rand(rng::Random.AbstractRNG, ::Type{T}, ::StdLogistic) where {T}
-    logit(rand(rng, T))
-end
+@inline rand_impl(ctx::GenContext, ::StdLogistic) = logit(rand(get_rng(ctx), get_precision(ctx)))
+@inline batched_rand_impl(ctx::GenContext, ::StdLogistic, sz::Dims) = logit.(_rand_bulk(ctx, sz))
 
 smf(::StdLogistic, x) = logistic(x)
 smf(::StdLogistic) = logistic
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index 316d0d40..044029d5 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -10,6 +10,8 @@ StdMeasure(::typeof(randn)) = StdNormal()
 
 @inline check_dof(::StdMeasure, ::StdMeasure) = nothing
 
+@inline massof(::StdMeasure) = static(1.0)
+
 @inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x
 
 @inline transport_to_std(::Type{S}, ::S, x) where {S<:StdMeasure} = x
diff --git a/src/standard/stdnormal.jl b/src/standard/stdnormal.jl
index d870d331..faa043dc 100644
--- a/src/standard/stdnormal.jl
+++ b/src/standard/stdnormal.jl
@@ -14,7 +14,8 @@ export StdNormal
 
 @inline getdof(::StdNormal) = static(1)
 
-@inline Base.rand(rng::Random.AbstractRNG, ::Type{T}, ::StdNormal) where {T} = randn(rng, T)
+@inline rand_impl(ctx::GenContext, ::StdNormal) = randn(get_rng(ctx), get_precision(ctx))
+@inline batched_rand_impl(ctx::GenContext, ::StdNormal, sz::Dims) = _randn_bulk(ctx, sz)
 
 Φ(z) = erfc(-z * invsqrt2) / 2
 Φinv(p) = -erfcinv(2 * p) * sqrt2
diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl
index f9b07bf1..7c66caec 100644
--- a/src/standard/stduniform.jl
+++ b/src/standard/stduniform.jl
@@ -12,7 +12,8 @@ end
 @inline logdensity_def(::StdUniform, x) = zero(x)
 @inline basemeasure(::StdUniform) = LebesgueBase()
 
-Base.rand(rng::Random.AbstractRNG, ::Type{T}, ::StdUniform) where {T} = rand(rng, T)
+@inline rand_impl(ctx::GenContext, ::StdUniform) = rand(get_rng(ctx), get_precision(ctx))
+@inline batched_rand_impl(ctx::GenContext, ::StdUniform, sz::Dims) = _rand_bulk(ctx, sz)
 
 massof(::StdUniform, s::Interval) = massof(Lebesgue(0.0 .. 1.0), s)
 
diff --git a/test/rand.jl b/test/rand.jl
new file mode 100644
index 00000000..c1899afe
--- /dev/null
+++ b/test/rand.jl
@@ -0,0 +1,109 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using Random, Statistics
+using StableRNGs: StableRNG
+using Static: static
+using ArraysOfArrays: flatview
+using AffineMaps: Add
+
+using MeasureBase
+using MeasureBase: GenContext
+using MeasureBase: StdUniform, StdExponential, StdLogistic, StdNormal, Dirac, Lebesgue
+using MeasureBase: weightedmeasure, superpose, mcombine, mbind, productmeasure, pushfwd, testvalue
+using MeasureBase: rand_impl, batched_rand_impl
+
+@testset "rand" begin
+    stblrng() = StableRNG(789990641)
+
+    @testset "generative contexts" begin
+        @test rand(stblrng(), StdNormal()) == rand(stblrng(), StdNormal())
+        @test rand(stblrng(), StdNormal()) == rand(GenContext{Float64}(stblrng()), StdNormal())
+        @test rand(stblrng(), StdNormal()^3) == rand(stblrng(), StdNormal()^3)
+        @test @inferred(rand(stblrng(), Float32, StdNormal())) isa Float32
+        @test @inferred(rand(Float32, StdNormal()^3)) isa Vector{Float32}
+        @test @inferred(rand(GenContext{Float32}(stblrng()), StdUniform()^(2, 3))) isa Matrix{Float32}
+        @test @inferred(rand(StdExponential())) isa Float64
+        @test_throws ArgumentError rand(Lebesgue())
+    end
+
+    @testset "layout of power variates" begin
+        x = rand(StdNormal()^(2, 3))
+        @test x isa Matrix{Float64} && size(x) == (2, 3)
+        xs = rand(StdNormal()^static(3))
+        @test xs isa AbstractVector{Float64} && length(xs) == 3
+        xn = rand((StdNormal()^2)^3)
+        @test xn isa AbstractVector && length(xn) == 3 && all(x -> length(x) == 2, xn)
+        @test size(flatview(xn)) == (2, 3)
+        @test logdensityof((StdNormal()^2)^3, xn) ≈ logdensityof(StdNormal()^6, vec(flatview(xn)))
+        @test batched_rand_impl(GenContext{Float64}(stblrng()), StdNormal()^2, (4, 5)) isa Array{Float64,3}
+        @test size(batched_rand_impl(GenContext{Float64}(stblrng()), (StdNormal()^2)^3, (4,))) == (2, 3, 4)
+    end
+
+    @testset "test values" begin
+        @test testvalue(StdNormal()) == 0
+        @test testvalue(Float32, StdUniform()) === 0.5f0
+        @test testvalue(StdExponential()^3) == ones(3)
+        @test testvalue((StdLogistic()^2)^2) == [zeros(2), zeros(2)]
+        @test testvalue(productmeasure((a = StdNormal(), b = StdUniform()^2))) == (a = 0.0, b = [0.5, 0.5])
+    end
+
+    @testset "distribution of variates" begin
+        n = 10^5
+        for (μ, m, v) in [
+            (StdUniform(), 0.5, 1 / 12),
+            (StdExponential(), 1.0, 1.0),
+            (StdLogistic(), 0.0, π^2 / 3),
+            (StdNormal(), 0.0, 1.0),
+            (weightedmeasure(0.3, StdNormal()), 0.0, 1.0),
+            (pushfwd(exp, StdNormal()), exp(0.5), (exp(1) - 1) * exp(1)),
+            (MeasureBase.Half(StdNormal()), sqrt(2 / π), 1 - 2 / π),
+            (SpikeMixture(Dirac(1.0), 0.25), 0.25, 0.25 * 0.75),
+        ]
+            X = rand(stblrng(), μ^n)
+            @test isapprox(mean(X), m, atol = 5 * sqrt(v / n) + 1e-3)
+            @test isapprox(var(X), v, rtol = 0.05)
+            xs = [rand_impl(GenContext{Float64}(stblrng()), μ) for _ in 1:20]
+            @test all(x -> insupport(μ, x) != false, xs)
+        end
+
+        mix = superpose(weightedmeasure(log(0.3), Dirac(0.0)), weightedmeasure(log(0.7), Dirac(1.0)))
+        @test isapprox(mean(rand(stblrng(), mix^n)), 0.7, atol = 0.01)
+        @test isapprox(mean([rand(mix) for _ in 1:n]), 0.7, atol = 0.01)
+        mixn = superpose(weightedmeasure(log(0.5), StdNormal()), weightedmeasure(log(0.5), pushfwd(Add(4.0), StdNormal())))
+        @test isapprox(mean(rand(stblrng(), mixn^n)), 2.0, atol = 0.02)
+    end
+
+    @testset "structural measures" begin
+        P = MeasureBase.ProductMeasure([weightedmeasure(log(i), StdNormal()) for i in 1:3])
+        @test @inferred(rand(stblrng(), P)) isa Vector{Float64}
+        XP = rand(stblrng(), P^100)
+        @test size(flatview(XP)) == (3, 100)
+        @test logdensities(P, XP) ≈ [logdensityof(P, x) for x in XP]
+
+        Pt = productmeasure((StdNormal(), StdUniform()^2))
+        xt = rand(stblrng(), Pt)
+        @test xt isa Tuple{Float64,Vector{Float64}}
+        Xt = rand(stblrng(), Pt^5)
+        @test Xt isa AbstractVector && length(Xt) == 5
+
+        mc = mcombine(vcat, StdNormal()^2, StdUniform()^3)
+        xc = rand(stblrng(), mc)
+        @test xc isa Vector{Float64} && length(xc) == 5 && all(0 .<= xc[3:5] .<= 1)
+        Xc = rand(stblrng(), mc^50)
+        @test size(flatview(Xc)) == (5, 50)
+        @test logdensities(mc, Xc) ≈ [logdensityof(mc, x) for x in Xc]
+
+        f_β(a) = StdNormal()^length(a)
+        μb = mbind(f_β, StdUniform()^2, vcat)
+        xb = rand(stblrng(), μb)
+        @test xb isa AbstractVector && length(xb) == 4
+        Xb = rand(stblrng(), μb^3)
+        @test Xb isa AbstractVector && length(Xb) == 3 && all(x -> length(x) == 4, Xb)
+
+        d = Dirac([1.0, 2.0])
+        @test rand(d^2) == [[1.0, 2.0], [1.0, 2.0]]
+        @test batched_rand_impl(GenContext{Float64}(stblrng()), d, (2,)) == [1.0 1.0; 2.0 2.0]
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index c95954a4..90cf12a2 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -35,6 +35,8 @@ include("combinators/implicitlymapped.jl")
 include("combinators/combined.jl")
 include("combinators/bind.jl")
 
+include("rand.jl")
+
 include("distributions/test_distributions.jl")
 
 include("test_docs.jl")

From 5dd262d14020502dc09e2eff64e08efd18c87512 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 01:56:22 +0200
Subject: [PATCH 090/122] Refuse abstract standard measure types as transport
 partners

Products over arrays of mixed standard measures had the abstract type
StdMeasure as preferred standard measure, which matched the identity
transports of standard measures and made transports between two such
products silent identities. Abstract standard measure types are no
preference anymore, so the marginal preferences promote at run time, and
they can't act as transport pivots. Array products of marginals with
mixed variate shapes transport marginal by marginal, the standard
distribution transports of the Distributions extension are no longer
ambiguous, and the extension transport tests cover uniform-normal pairs
as their names suggest.

Created by generative AI.
---
 .../standard_dist.jl                          |  4 ++--
 src/combinators/product.jl                    | 10 +++++++-
 src/standard/stdmeasure.jl                    | 16 ++++++++++---
 src/transport.jl                              | 23 ++++++++++++++++---
 test/distributions/test_transport.jl          |  5 ++--
 test/transport.jl                             | 19 +++++++++++++++
 6 files changed, 65 insertions(+), 12 deletions(-)

diff --git a/ext/MeasureBaseDistributionsExt/standard_dist.jl b/ext/MeasureBaseDistributionsExt/standard_dist.jl
index 8e8963ad..a51280b3 100644
--- a/ext/MeasureBaseDistributionsExt/standard_dist.jl
+++ b/ext/MeasureBaseDistributionsExt/standard_dist.jl
@@ -41,8 +41,8 @@ for (A, B) in [
 ]
     @eval begin
         @inline MeasureBase.preferred_stdmeasure(::Type{<:StandardDist{$A}}) = $B
-        @inline MeasureBase.transport_to_std(::Type{$B}, ::StandardDist{$A,0}, x) = x
-        @inline MeasureBase.transport_from_std(::Type{$B}, ::StandardDist{$A,0}, z) = z
+        @inline MeasureBase.transport_to_std(::Type{$B}, ::StandardDist{$A,0}, x::Number) = x
+        @inline MeasureBase.transport_from_std(::Type{$B}, ::StandardDist{$A,0}, z::Number) = z
         @inline MeasureBase.transport_to_std(::Type{$B}, ::StandardDist{$A}, x::AbstractArray) = vec(x)
         @inline MeasureBase.transport_from_std(::Type{$B}, d::StandardDist{$A}, z::AbstractVector) = reshape(z, size(d))
 
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index fc5c37fa..bfb94e08 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -337,9 +337,17 @@ function transport_to_std(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, x:
     transport_to_std(S, productmeasure(values(marginals(μ))), values(x))
 end
 
-function transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray}, x::AbstractArray) where {S<:StdMeasure}
+function transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray) where {S<:StdMeasure,M}
+    _array_product_to_std(S, μ, x, Val(isconcretetype(M)))
+end
+function _array_product_to_std(::Type{S}, μ, x::AbstractArray, ::Val{true}) where {S}
     _flat_std_of(broadcast(_ToStd{S}(), marginals(μ), x))
 end
+# Marginals of mixed types may have standard variates of mixed shapes:
+function _array_product_to_std(::Type{S}, μ, x::AbstractArray, ::Val{false}) where {S}
+    zs = [_as_stdstream(transport_to_std(S, m, xi)) for (m, xi) in zip(marginals(μ), x)]
+    isempty(zs) ? SVector{0,Bool}() : reduce(vcat, zs)
+end
 
 function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:Tuple}, z::AbstractVector) where {S<:StdMeasure}
     _marginals_from_std_with_rest(S, marginals(μ), z)
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index 044029d5..ea767146 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -14,8 +14,16 @@ StdMeasure(::typeof(randn)) = StdNormal()
 
 @inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x
 
-@inline transport_to_std(::Type{S}, ::S, x) where {S<:StdMeasure} = x
-@inline transport_from_std(::Type{S}, ::S, z) where {S<:StdMeasure} = z
+@inline transport_to_std(::Type{S}, ::S, x) where {S<:StdMeasure} = _std_identity(S, x)
+@inline transport_from_std(::Type{S}, ::S, z) where {S<:StdMeasure} = _std_identity(S, z)
+
+# Only concrete standard measure types identify a transport partner:
+@inline function _std_identity(::Type{S}, x) where {S<:StdMeasure}
+    isconcretetype(S) || _throw_abstract_std(S)
+    return x
+end
+@noinline _throw_abstract_std(::Type{S}) where {S} =
+    throw(ArgumentError("$(S) is not a concrete standard measure type"))
 
 
 """
@@ -56,7 +64,9 @@ function preferred_stdmeasure end
 @inline preferred_stdmeasure(μ) = preferred_stdmeasure(typeof(μ))
 @inline preferred_stdmeasure(::Type{MU}) where {MU} = NoStdTransport{MU}
 
-@inline preferred_stdmeasure(::Type{MU}) where {MU<:StdMeasure} = MU
+@inline function preferred_stdmeasure(::Type{MU}) where {MU<:StdMeasure}
+    isconcretetype(MU) ? MU : NoStdTransport{MU}
+end
 
 """
     MeasureBase.promote_stdmeasure(A::Type, B::Type)::Type
diff --git a/src/transport.jl b/src/transport.jl
index 0c88e0b6..d6526012 100644
--- a/src/transport.jl
+++ b/src/transport.jl
@@ -21,8 +21,9 @@ transport_to(StdNormal, μ)
 transport_to(ν, StdNormal)
 ```
 
-An instance of the standard measure itself or a power of it (depending on
-[`getdof(μ)`](@ref) resp. `ν`) is chosen as the transport partner then.
+The transport partner is then an instance of the standard measure for
+measures with scalar variates, and a power of it with as many elements as
+the measure has degrees of freedom otherwise.
 
 # Extended help
 
@@ -143,7 +144,10 @@ end
 @inline function _transport_pivot(ν, μ)
     _concrete_pivot(promote_stdmeasure(preferred_stdmeasure(ν), preferred_stdmeasure(μ)), ν, μ)
 end
-@inline _concrete_pivot(::Type{S}, ν, μ) where {S<:StdMeasure} = S
+@inline function _concrete_pivot(::Type{S}, ν, μ) where {S<:StdMeasure}
+    isconcretetype(S) || _throw_abstract_std(S)
+    return S
+end
 @inline _concrete_pivot(::Type{AnyStdMeasure}, ν, μ) = StdUniform
 function _concrete_pivot(::Type{<:NoStdTransport{MU}}, ν, μ) where {MU}
     throw(ArgumentError("No transport between measures of type $(nameof(typeof(ν))) and $(nameof(typeof(μ))), measures of type $(nameof(MU)) have no transport via standard measures"))
@@ -178,6 +182,9 @@ end
 function _to_std_via(::Type{S}, ::Type{S}, μ, x) where {S<:StdMeasure}
     throw(ArgumentError("Transport to $(nameof(S)) is not implemented for measures of type $(nameof(typeof(μ)))"))
 end
+function _to_std_via(::Type{S}, ::Type{AnyStdMeasure}, μ, x) where {S<:StdMeasure}
+    throw(ArgumentError("Transport to standard measures is not implemented for measures of type $(nameof(typeof(μ)))"))
+end
 function _to_std_via(::Type{S}, ::Type, μ, x) where {S<:StdMeasure}
     throw(ArgumentError("Measures of type $(nameof(typeof(μ))) have no transport via standard measures"))
 end
@@ -201,6 +208,9 @@ end
 function _from_std_via(::Type{S}, ::Type{S}, μ, z) where {S<:StdMeasure}
     throw(ArgumentError("Transport from $(nameof(S)) is not implemented for measures of type $(nameof(typeof(μ)))"))
 end
+function _from_std_via(::Type{S}, ::Type{AnyStdMeasure}, μ, z) where {S<:StdMeasure}
+    throw(ArgumentError("Transport from standard measures is not implemented for measures of type $(nameof(typeof(μ)))"))
+end
 function _from_std_via(::Type{S}, ::Type, μ, z) where {S<:StdMeasure}
     throw(ArgumentError("Measures of type $(nameof(typeof(μ))) have no transport via standard measures"))
 end
@@ -216,6 +226,13 @@ Returns a tuple `(z, x_μ, x_rest)` of the flat vector `z` of standard
 variates, the variate `x_μ` of `μ` consumed from the stream and the
 unconsumed rest of the stream. See
 [`MeasureBase.logdensityof_with_rest`](@ref) for the stream conventions.
+
+The default implementation consumes a variate of the size given by
+[`MeasureBase.mspace_flatsize`](@ref) or
+[`MeasureBase.some_mspace_elsize`](@ref). Measure types whose variates are
+composed of the variates of other measures implement
+`transport_to_std_with_rest` instead of
+[`MeasureBase.transport_to_std`](@ref).
 """
 function transport_to_std_with_rest end
 
diff --git a/test/distributions/test_transport.jl b/test/distributions/test_transport.jl
index c8d719b2..bdeb33c9 100644
--- a/test/distributions/test_transport.jl
+++ b/test/distributions/test_transport.jl
@@ -7,11 +7,10 @@ using InverseFunctions, ChangesOfVariables
 using Distributions, ArraysOfArrays
 using ArraysOfArrays: sliced, flatview
 using StableRNGs
-using LogExpFunctions: logit
 import ForwardDiff, Zygote
 import PDMats
 
-using MeasureBase: transport_to, transport_def
+using MeasureBase: transport_to
 using MeasureBase: StdUniform, StdNormal, StdExponential, StdLogistic
 using .MeasureBaseDistributionsExt: _trafo_logcdf, _trafo_logccdf, _trafo_quantile, _trafo_cquantile
 
@@ -49,7 +48,7 @@ include("getjacobian.jl")
 
     @testset "transforms-tests" begin
         stduvuni = StandardDist{Uniform}()
-        stduvnorm = StandardDist{Uniform}()
+        stduvnorm = StandardDist{Normal}()
 
         uniform1 = Uniform(-5.0, -0.01)
         uniform2 = Uniform(0.01, 5.0)
diff --git a/test/transport.jl b/test/transport.jl
index b239de90..c57c12d8 100644
--- a/test/transport.jl
+++ b/test/transport.jl
@@ -130,6 +130,25 @@ using JLArrays
         @test transport_to(StdUniform(), μ)(0.0) ≈ 0
     end
 
+    @testset "array products of mixed standard measures" begin
+        src = productmeasure([StdNormal(), StdExponential()])
+        trg = productmeasure([StdUniform(), StdLogistic()])
+        @test MeasureBase.preferred_stdmeasure(src) === StdNormal
+        f = transport_to(trg, src)
+        x = [0.5, 1.0]
+        y = f(x)
+        @test y ≈ [transport_to(StdUniform(), StdNormal())(0.5), transport_to(StdLogistic(), StdExponential())(1.0)]
+        @test inverse(f)(y) ≈ x
+        @test_throws ArgumentError transport_to_std(MeasureBase.StdMeasure, StdNormal(), 0.5)
+
+        pm = productmeasure(AbstractMeasure[StdNormal(), StdNormal()^2, Dirac(1.0)])
+        xm = [0.5, [0.1, 0.2], 1.0]
+        z = transport_to(StdUniform()^3, pm)(xm)
+        @test z isa AbstractVector{<:Real} && length(z) == 3
+        xm_reco = transport_to(pm, StdUniform()^3)(z)
+        @test xm_reco[1] ≈ xm[1] && xm_reco[2] ≈ xm[2] && xm_reco[3] == 1.0
+    end
+
     @testset "measures without standard transport" begin
         μ = restrict(x -> x > 0, StdNormal())
         @test_throws ArgumentError transport_to(StdUniform(), μ)(0.5)

From d97f084bc89af95e015a88259cfcd320dfdc180c Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 02:36:11 +0200
Subject: [PATCH 091/122] Fix batched generation and transport of mixtures,
 powers and products

Constant-RNG contexts now really fill constant bulks (the specialization
had lost to the generic method), so test values of combined measures
work again. The broadcast hook of transport functions no longer re-enters
itself through GPU array maps, keeps static arrays static, materializes
fused arguments and re-nests variates of nested powers. Batches of spike
mixtures and superpositions with array variates align their masks with
the flat batch, variates generated via standard transports keep the
context precision, fused array-product transports require marginals
with one degree of freedom, powers have the right mass and standard
measures report as normalized.

Created by generative AI.
---
 src/combinators/power.jl              |  2 +-
 src/combinators/product.jl            | 21 ++++++++++------
 src/combinators/spikemixture.jl       |  7 +++++-
 src/combinators/superpose.jl          |  7 +++++-
 src/combinators/transformedmeasure.jl |  5 +---
 src/mass-interface.jl                 |  2 +-
 src/rand.jl                           | 20 ++++++++++-----
 src/transport-batched.jl              | 35 +++++++++++++++++++++------
 src/transport.jl                      |  9 ++++---
 test/distributions/test_transport.jl  | 11 +++++++++
 test/rand.jl                          | 15 +++++++++++-
 test/transport.jl                     | 17 +++++++++++--
 12 files changed, 117 insertions(+), 34 deletions(-)

diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 026a65c5..b48f38f5 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -180,7 +180,7 @@ end
 
 checked_arg(μ::PowerMeasure, x::Any) = _throw_size_mismatch()
 
-massof(m::PowerMeasure) = massof(m.parent)^prod(m.axes)
+massof(m::PowerMeasure) = massof(m.parent)^dynamic(size2length(pwr_size(m)))
 
 
 # Transport: the standard variate of a power is the flat vector of the
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index bfb94e08..1d46e93a 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -31,7 +31,7 @@ basemeasure(d::AbstractProductMeasure) = productmeasure(map(basemeasure, margina
 rand_impl(ctx::GenContext, d::AbstractProductMeasure) = map(Base.Fix1(_marginal_rand, ctx), marginals(d))
 
 @inline _marginal_rand(ctx::GenContext, m::AbstractMeasure) = rand_impl(ctx, m)
-@inline _marginal_rand(ctx::GenContext, d) = rand(get_rng(ctx), d)
+@inline _marginal_rand(ctx::GenContext, d) = convert_realtype(get_precision(ctx), rand(get_rng(ctx), d))
 
 for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)]
     @eval @inline function $head(d::AbstractProductMeasure, x)
@@ -287,6 +287,7 @@ fast_dof(d::AbstractProductMeasure) = _sum_dofs(fast_dof, marginals(d))
 # are summed dynamically (also on GPU arrays):
 @inline _sum_dofs(f, mar) = sum(f, mar)
 @inline _sum_dofs(f, mar::AbstractArray) = mapreduce(_dynamic_dof ∘ f, +, mar; init = 0)
+@inline _sum_dofs(f, mar::StaticArray) = mapreduce(f, +, mar; init = static(0))
 @inline _dynamic_dof(n::IntegerLike) = dynamic(n)
 @inline _dynamic_dof(nodof::AbstractNoDOF) = nodof
 
@@ -417,24 +418,30 @@ end
 # Batched transport of array products with scalar-variate marginals in one
 # broadcast, the marginals align with the leading dimension of the batch:
 
+# Marginals of concrete type with scalar variates and a standard transport
+# have one degree of freedom each, so the batch aligns with the marginals:
+@inline function _fused_marginals(::Type{M}) where {M}
+    Val(isconcretetype(M) && mspace_flatsize(M) === () && preferred_stdmeasure(M) isa Type{<:StdMeasure})
+end
+
 function batched_transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, X::AbstractArray) where {S<:StdMeasure,M}
-    _array_product_batched_to_std(S, μ, X, mspace_flatsize(M), Val(isconcretetype(M)))
+    _array_product_batched_to_std(S, μ, X, _fused_marginals(M))
 end
-function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Tuple{}, ::Val{true}) where {S}
+function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Val{true}) where {S}
     _check_flatsize(X, maybestatic_size(marginals(μ)))
     _as_stream_batch(broadcast(_ToStd{S}(), marginals(μ), X), maybestatic_size(marginals(μ)))
 end
-function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Any, ::Val) where {S}
+function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Val{false}) where {S}
     _batched_to_std(S, μ, X, mspace_flatsize(μ))
 end
 
 function batched_transport_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, Z::AbstractArray) where {S<:StdMeasure,M}
-    _array_product_batched_from_std(S, μ, Z, mspace_flatsize(M), Val(isconcretetype(M)))
+    _array_product_batched_from_std(S, μ, Z, _fused_marginals(M))
 end
-function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Tuple{}, ::Val{true}) where {S}
+function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{true}) where {S}
     mar = marginals(μ)
     broadcast(_FromStd{S}(), mar, reshape(Z, (map(dynamic, maybestatic_size(mar))..., Base.tail(size(Z))...)))
 end
-function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Any, ::Val) where {S}
+function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{false}) where {S}
     _batched_from_std(S, μ, Z, mspace_flatsize(μ))
 end
diff --git a/src/combinators/spikemixture.jl b/src/combinators/spikemixture.jl
index 440c67d5..58d8ad19 100644
--- a/src/combinators/spikemixture.jl
+++ b/src/combinators/spikemixture.jl
@@ -38,9 +38,14 @@ function rand_impl(ctx::GenContext, μ::SpikeMixture)
 end
 
 function batched_rand_impl(ctx::GenContext, μ::SpikeMixture, sz::Dims)
+    _spike_batched_rand(ctx, μ, sz, mspace_flatsize(μ.m))
+end
+function _spike_batched_rand(ctx::GenContext, μ::SpikeMixture, sz::Dims, sz_flat::SizeLike)
     X = batched_rand_impl(ctx, μ.m, sz)
-    return ifelse.(_rand_bulk(ctx, sz) .< μ.w, X, zero(eltype(X)))
+    return ifelse.(_batch_mask(_rand_bulk(ctx, sz) .< μ.w, sz_flat), X, zero(eltype(X)))
 end
+_spike_batched_rand(ctx::GenContext, μ::SpikeMixture, sz::Dims, ::NoMSpaceElementSize) =
+    _batched_rand_pointwise(ctx, μ, sz)
 
 testvalue(::Type{T}, μ::SpikeMixture) where {T} = zero(T)
 
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index 1c81886e..3ca5d554 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -149,9 +149,12 @@ end
 # Batches of superpositions draw a batch from each component and select
 # by mass, branch-free:
 function batched_rand_impl(ctx::GenContext, μ::SuperpositionMeasure, sz::Dims)
+    _superpose_batched_rand(ctx, μ, sz, mspace_flatsize(μ))
+end
+function _superpose_batched_rand(ctx::GenContext, μ::SuperpositionMeasure, sz::Dims, sz_flat::SizeLike)
     components = values(μ.components)
     masses, total = _component_masses(μ)
-    thresholds = _rand_bulk(ctx, sz) .* total
+    thresholds = _batch_mask(_rand_bulk(ctx, sz) .* total, sz_flat)
     X = batched_rand_impl(ctx, first(components), sz)
     csum = first(masses)
     for (mass, c) in Iterators.drop(zip(masses, components), 1)
@@ -160,6 +163,8 @@ function batched_rand_impl(ctx::GenContext, μ::SuperpositionMeasure, sz::Dims)
     end
     return X
 end
+_superpose_batched_rand(ctx::GenContext, μ::SuperpositionMeasure, sz::Dims, ::NoMSpaceElementSize) =
+    _batched_rand_pointwise(ctx, μ, sz)
 
 @inline function insupport(d::SuperpositionMeasure, x)
     mapreduce(c -> _insupport_mask(insupport(c, x)), |, values(d.components))
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 7db728e8..b8d5d0c4 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -245,10 +245,7 @@ end
 @inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, ::Tuple{})
     broadcast(ν.f, batched_rand_impl(ctx, ν.origin, sz))
 end
-function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, sz_orig::SizeLike)
-    stacked(map(ν.f, sliced(batched_rand_impl(ctx, ν.origin, sz), Val(length(sz_orig)))))
-end
-@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, ::NoMSpaceElementSize)
+@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, ::Any)
     _batched_rand_pointwise(ctx, ν, sz)
 end
 
diff --git a/src/mass-interface.jl b/src/mass-interface.jl
index 548a941c..0807b6e6 100644
--- a/src/mass-interface.jl
+++ b/src/mass-interface.jl
@@ -107,7 +107,7 @@ Check whether `norm(x, p) == 1`.
 """
 isnormalized(x, p::Real = 2) = isone(norm(x, p))
 
-isone(::AbstractUnknownMass) = false
+Base.isone(::AbstractUnknownMass) = false
 
 function massof(m, s)
     _default_massof_impl(m, s, rootmeasure(m))
diff --git a/src/rand.jl b/src/rand.jl
index cf8a8a5e..5168a087 100644
--- a/src/rand.jl
+++ b/src/rand.jl
@@ -12,7 +12,8 @@ Generate a random variate of `μ`.
 The generative context `ctx` (see `HeterogeneousComputing.GenContext`)
 determines the random number generator, the numerical precision (`Float64`
 by default) and the compute unit that array-valued variates are generated
-on. The variates of powers of measures are generated in one batch.
+on. The variates of powers of measures with a known flat variate size are
+generated in one batch.
 
 Measure types should specialize [`MeasureBase.rand_impl`](@ref) and
 [`MeasureBase.batched_rand_impl`](@ref) instead of `rand`.
@@ -45,10 +46,10 @@ function rand_impl(ctx::GenContext, μ)
 end
 
 @inline function _rand_via_std(ctx::GenContext, μ, ::Type{S}, ::IntegerLike, ::Tuple{}) where {S<:StdMeasure}
-    transport_from_std(S, μ, rand_impl(ctx, S()))
+    convert_realtype(get_precision(ctx), transport_from_std(S, μ, rand_impl(ctx, S())))
 end
 @inline function _rand_via_std(ctx::GenContext, μ, ::Type{S}, n::IntegerLike, ::Any) where {S<:StdMeasure}
-    transport_from_std(S, μ, _rand_std(ctx, S, (dynamic(n),)))
+    convert_realtype(get_precision(ctx), transport_from_std(S, μ, _rand_std(ctx, S, (dynamic(n),))))
 end
 @inline function _rand_via_std(ctx::GenContext, μ, ::Type{AnyStdMeasure}, n::IntegerLike, sz)
     _rand_via_std(ctx, μ, StdUniform, n, sz)
@@ -63,7 +64,8 @@ end
 
 Generate a batch of random variates of `μ` of batch size `sz` in flat
 form, an array of size `(flat variate dims..., sz...)` (see
-[`MeasureBase.mspace_flatsize`](@ref)).
+[`MeasureBase.mspace_flatsize`](@ref)). Measures with variates of
+unknown flat size only support batches of scalar variates.
 
 The default implementation draws a batch of variates of the preferred
 standard measure of `μ` and transports it to `μ`, or generates the
@@ -76,7 +78,7 @@ function batched_rand_impl(ctx::GenContext, μ, sz::Dims)
 end
 
 function _batched_rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{S}, n::IntegerLike, ::SizeLike) where {S<:StdMeasure}
-    batched_transport_from_std(S, μ, _rand_std(ctx, S, (dynamic(n), sz...)))
+    convert_realtype(get_precision(ctx), batched_transport_from_std(S, μ, _rand_std(ctx, S, (dynamic(n), sz...))))
 end
 @inline function _batched_rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{AnyStdMeasure}, n::IntegerLike, sz_flat::SizeLike)
     _batched_rand_via_std(ctx, μ, sz, StdUniform, n, sz_flat)
@@ -114,12 +116,18 @@ end
 @inline _randexp_bulk(ctx::GenContext, sz::Dims, ::AbstractComputeUnit) = -log1p.(-_rand_bulk(ctx, sz))
 
 # Test values use a constant RNG, which only draws single values:
-const _ConstantContext = GenContext{<:Any,<:Any,ConstantRNG}
+const _ConstantContext = GenContext{<:AbstractFloat,<:AbstractComputeUnit,ConstantRNG}
 @inline _rand_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, rand(ConstantRNG(), get_precision(ctx)), sz)
 @inline _randn_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, randn(ConstantRNG(), get_precision(ctx)), sz)
 @inline _randexp_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, randexp(ConstantRNG(), get_precision(ctx)), sz)
 @inline _const_bulk(ctx::GenContext, x, sz::Dims) = fill!(allocate_array(ctx, typeof(x), sz), x)
 
+# A mask over the batch dimensions, aligned with a flat batch of variates:
+@inline _batch_mask(mask::AbstractArray, ::Tuple{}) = mask
+@inline function _batch_mask(mask::AbstractArray, sz_flat::SizeLike)
+    reshape(mask, (ntuple(_ -> 1, Val(length(sz_flat)))..., size(mask)...))
+end
+
 # A batch of copies of a constant variate:
 function _const_batch(ctx::GenContext, x, sz::Dims)
     X = allocate_array(ctx, eltype(x), (size(x)..., sz...))
diff --git a/src/transport-batched.jl b/src/transport-batched.jl
index d14ea7ca..5a8cedcf 100644
--- a/src/transport-batched.jl
+++ b/src/transport-batched.jl
@@ -12,7 +12,7 @@ of variates of the standard measure type `S`.
 
 The default implementation broadcasts the point transport for measures
 with scalar variates and maps it over the variate slices of `X`
-otherwise.
+otherwise, in a host loop.
 """
 function batched_transport_to_std end
 
@@ -44,6 +44,10 @@ end
 Batched form of [`MeasureBase.transport_from_std`](@ref): transports the
 batch `Z` of variates of the standard measure type `S`, of size
 `(getdof(μ), batch dims...)`, to a flat batch of variates of `μ`.
+
+The default implementation broadcasts the point transport for measures
+with scalar variates and maps it over the columns of `Z` otherwise, in a
+host loop.
 """
 function batched_transport_from_std end
 
@@ -72,7 +76,8 @@ Batched form of [`MeasureBase.transport_to_std_with_rest`](@ref) for a
 batch `X` of flat vector streams (first dimension along the streams).
 
 Returns a tuple `(Z, X_μ, X_rest)` of the batch of standard variates, the
-rows consumed from the streams and the unconsumed rest of the streams.
+batch of variates of `μ` consumed from the streams and the unconsumed rest
+of the streams.
 """
 function batched_transport_to_std_with_rest end
 
@@ -139,18 +144,34 @@ end
 
 # Broadcasting a transport function over an array of variates with flat
 # storage, or over the flat storage of a batch, transports the batch as a
-# whole. Variates of the target measure come out in their flat form.
+# whole. Fused broadcast arguments are materialized first, static arrays
+# are transported point by point.
 function Broadcast.broadcasted(f::TransportFunction, X::AbstractArray)
     _broadcast_transport(f, X, _flat_storage(X), mspace_flatsize(f.μ), mspace_flatsize(f.ν))
 end
 
+function Broadcast.broadcasted(f::TransportFunction, bc::Broadcast.Broadcasted)
+    Broadcast.broadcasted(f, Broadcast.materialize(bc))
+end
+
+Broadcast.broadcasted(f::TransportFunction, X::StaticArray) = map(_Pointwise(f), X)
+
 function _broadcast_transport(f::TransportFunction, X, X_flat::AbstractArray, sz_μ::SizeLike, sz_ν::SizeLike)
     _check_flatsize(X_flat, sz_μ)
     Y_flat = batched_transport_def(f.ν, f.μ, X_flat)
-    return _batch_variates(Y_flat, sz_ν)
+    return _batch_variates(Y_flat, f.ν)
 end
 
-_broadcast_transport(f::TransportFunction, X, ::Any, ::Any, ::Any) = map(f, X)
+_broadcast_transport(f::TransportFunction, X, ::Any, ::Any, ::Any) = map(_Pointwise(f), X)
+
+# Prevents re-entering the broadcast hook from `map` implementations that
+# broadcast (e.g. GPU arrays):
+struct _Pointwise{F} <: Function
+    f::F
+end
+@inline (p::_Pointwise)(x) = p.f(x)
 
-@inline _batch_variates(Y::AbstractArray, ::Tuple{}) = Y
-@inline _batch_variates(Y::AbstractArray, sz::SizeLike) = sliced(Y, Val(length(sz)))
+# The batch of variates in the layout of the target measure over the flat
+# result, nested powers included:
+@inline _batch_variates(Y::AbstractArray, ν) = _nest_leaf(Y, mspace_flatsize(ν))
+@inline _batch_variates(Y::AbstractArray, ν::PowerMeasure) = sliced(_pwr_variate(ν, Y), Val(length(pwr_axes(ν))))
diff --git a/src/transport.jl b/src/transport.jl
index d6526012..af9116b7 100644
--- a/src/transport.jl
+++ b/src/transport.jl
@@ -12,9 +12,12 @@ to `rand(ν)`. `f` supports `InverseFunctions.inverse` and
 
 Measures are transported via standard measures: `x` is transported to the
 standard measure type that the preferences of `ν` and `μ` promote to (see
-[`MeasureBase.preferred_stdmeasure`](@ref)) and from there to `ν`. A
-standard measure type like `StdUniform` or `StdNormal` may also be used
-directly as the source or target:
+[`MeasureBase.preferred_stdmeasure`](@ref)) and from there to `ν`.
+Broadcasting `f` over an array of variates with flat storage (see
+[`MeasureBase.mspace_flatsize`](@ref)), or over the flat storage of a
+batch of variates, transports the whole batch at once. A standard measure
+type like `StdUniform` or `StdNormal` may also be used directly as the
+source or target:
 
 ```julia
 transport_to(StdNormal, μ)
diff --git a/test/distributions/test_transport.jl b/test/distributions/test_transport.jl
index bdeb33c9..da14b275 100644
--- a/test/distributions/test_transport.jl
+++ b/test/distributions/test_transport.jl
@@ -233,6 +233,17 @@ include("getjacobian.jl")
         h = transport_to(StdNormal()^3, asmeasure(pd))
         Xp = rand(StableRNG(789990641), pd, 5)
         @test stack(h.(sliced(Xp, Val(1)))) ≈ stack(map(h, eachcol(Xp)))
+        pn = product_distribution([Normal(1.0, 2.0), Normal(0.0, 3.0), Normal(2.0, 1.0)])
+        mn = asmeasure(pn)
+        @test MeasureBase.mspace_flatsize(mn) == (3,)
+        hn = transport_to(StdUniform()^3, mn)
+        Xn = rand(StableRNG(789990641), pn, 4)
+        Yn = hn.(sliced(Xn, Val(1)))
+        @test flatview(Yn) ≈ stack(map(hn, eachcol(Xn)))
+        @test flatview(inverse(hn).(Yn)) ≈ Xn
+        @test eltype(rand(StableRNG(1), Float32, mn)) == Float32
+        @test eltype(flatview(rand(StableRNG(1), Float32, mn^3))) == Float32
+        @test eltype(rand(StableRNG(1), Float32, asmeasure(pd))) == Float32
     end
 
     @testset "MvNormal covariance representations" begin
diff --git a/test/rand.jl b/test/rand.jl
index c1899afe..867c6276 100644
--- a/test/rand.jl
+++ b/test/rand.jl
@@ -12,7 +12,7 @@ using MeasureBase
 using MeasureBase: GenContext
 using MeasureBase: StdUniform, StdExponential, StdLogistic, StdNormal, Dirac, Lebesgue
 using MeasureBase: weightedmeasure, superpose, mcombine, mbind, productmeasure, pushfwd, testvalue
-using MeasureBase: rand_impl, batched_rand_impl
+using MeasureBase: rand_impl, batched_rand_impl, massof, isnormalized
 
 @testset "rand" begin
     stblrng() = StableRNG(789990641)
@@ -42,6 +42,9 @@ using MeasureBase: rand_impl, batched_rand_impl
     end
 
     @testset "test values" begin
+        @test testvalue(mcombine(vcat, StdNormal()^2, StdUniform()^3)) == [0.0, 0.0, 0.5, 0.5, 0.5]
+        @test rand(MeasureBase.ConstantRNG(), Float64, StdUniform()^3) == fill(0.5, 3)
+        @test rand(MeasureBase.ConstantRNG(), Float32, StdLogistic()^2) == zeros(Float32, 2)
         @test testvalue(StdNormal()) == 0
         @test testvalue(Float32, StdUniform()) === 0.5f0
         @test testvalue(StdExponential()^3) == ones(3)
@@ -102,6 +105,16 @@ using MeasureBase: rand_impl, batched_rand_impl
         Xb = rand(stblrng(), μb^3)
         @test Xb isa AbstractVector && length(Xb) == 3 && all(x -> length(x) == 4, Xb)
 
+        ctx = GenContext{Float64}(stblrng())
+        spd = superpose(weightedmeasure(log(0.5), Dirac([1.0, 2.0])), weightedmeasure(log(0.5), Dirac([3.0, 4.0])))
+        Xspd = batched_rand_impl(ctx, spd, (6,))
+        @test size(Xspd) == (2, 6) && all(c -> c == [1.0, 2.0] || c == [3.0, 4.0], eachcol(Xspd))
+        Xsm = batched_rand_impl(ctx, SpikeMixture(StdNormal()^3, 0.5), (4,))
+        @test size(Xsm) == (3, 4) && all(c -> all(iszero, c) || !any(iszero, c), eachcol(Xsm))
+        @test massof(StdNormal()) == 1 && massof(StdUniform()^3) == 1 && massof(weightedmeasure(log(2.0), StdNormal()^2)) ≈ 2
+        @test isnormalized(StdNormal()) && isnormalized(StdExponential()^(2, 2)) && !isnormalized(2.0 * StdNormal())
+        @test !isnormalized(Lebesgue())
+
         d = Dirac([1.0, 2.0])
         @test rand(d^2) == [[1.0, 2.0], [1.0, 2.0]]
         @test batched_rand_impl(GenContext{Float64}(stblrng()), d, (2,)) == [1.0 1.0; 2.0 2.0]
diff --git a/test/transport.jl b/test/transport.jl
index c57c12d8..ea3c8614 100644
--- a/test/transport.jl
+++ b/test/transport.jl
@@ -9,7 +9,7 @@ using MeasureBase: weightedmeasure, mcombine
 using StaticArrays: SVector
 using Static: static
 using LogExpFunctions: logit
-using ArraysOfArrays: sliced, flatview
+using ArraysOfArrays: sliced, flatview, fused
 using JLArrays
 
 @testset "transport_to" begin
@@ -175,7 +175,18 @@ using JLArrays
         h = transport_to(StdUniform()^(2, 3), (StdNormal()^2)^3)
         Xh = randn(2, 3, 4)
         @test flatview(h.(Xh)) ≈ stack([h(Xh[:, :, i]) for i in 1:4])
-        @test flatview(inverse(h).(h.(Xh))) ≈ Xh
+        @test flatview(fused(inverse(h).(h.(Xh)))) ≈ Xh
+        Yh = h.(Xh)
+        Xh_reco = inverse(h).(Yh)
+        @test Xh_reco[2] == inverse(h)(Yh[2])
+        @test Xh_reco[2] isa AbstractVector && length(Xh_reco[2]) == 3 && Xh_reco[2][1] isa AbstractVector
+
+        X3 = randn(3, 4, 5)
+        Y3 = g.(X3)
+        @test size(Y3) == (4, 5) && size(flatview(Y3)) == (3, 4, 5)
+        @test Y3[2, 3] ≈ g(X3[:, 2, 3])
+        @test f.(SVector(0.3, 0.6, 0.9)) isa SVector{3,Float64}
+        @test g.(Xn .+ 0.0) == g.(Xn)
 
         P = MeasureBase.ProductMeasure([weightedmeasure(log(i), StdNormal()) for i in 1:3])
         p = transport_to(StdUniform()^3, P)
@@ -211,6 +222,8 @@ using JLArrays
         Pj = MeasureBase.ProductMeasure(JLArray([weightedmeasure(log(i), StdNormal()) for i in 1:3]))
         pj = transport_to(StdUniform()^3, Pj)
         @test Array(flatview(pj.(sliced(JLArray(Xp), Val(1))))) ≈ flatview(Yp)
+        Yej = pf.(JLArray(Xe))
+        @test Yej isa JLArray && Array(Yej) ≈ pf.(Xe)
     end
 
     @testset "transport for products" begin

From f35f91429f72fe0e2d765cd9d834efb5980bd8ad Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 09:56:11 +0200
Subject: [PATCH 092/122] Make batched density kernels the primary extension
 point

Density evaluation is batched first: batched_logdensityof_impl and
batched_logdensity_def act on flat batches with the variate dimensions
leading, zero batch dimensions being a single variate, and point
evaluation of array variates goes through them. Plain arrays of numbers
are flat storage by convention, so the kernels rely on the variate rank
of their measure, declared via mspace_ndims or derived from the flat
variate size, instead of routing on sizes. Stream consumption takes a
multiplicity, powers consume their size times the base's variates and
sum, chains with variates of value-dependent size are evaluated stream by
stream by the outermost combinator, and pushforwards learn their variate
size at construction from a test value of the origin.

Created by generative AI.
---
 .../distribution_measure.jl                   |   2 +
 src/MeasureBase.jl                            |   2 +-
 src/combinators/bind.jl                       |  19 +-
 src/combinators/combined.jl                   |  69 ++-
 src/combinators/half.jl                       |   2 +
 src/combinators/power.jl                      |  79 +++-
 src/combinators/product.jl                    | 110 ++++-
 src/combinators/restricted.jl                 |   2 +
 src/combinators/superpose.jl                  |  11 +
 src/combinators/transformedmeasure.jl         |  38 +-
 src/combinators/weighted.jl                   |   9 +-
 src/density-batched.jl                        | 429 +++++++++---------
 src/density-core.jl                           |  38 +-
 src/density.jl                                |   2 +
 src/mspace.jl                                 |  27 ++
 src/primitives/dirac.jl                       |  14 +
 src/transport-batched.jl                      |   8 +-
 test/combinators/combined.jl                  |   3 +-
 test/logdensities.jl                          |   8 +-
 19 files changed, 572 insertions(+), 300 deletions(-)

diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index 652826ad..9bc59610 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -64,6 +64,8 @@ end
 @inline MeasureBase.mspace_elsize(m::DistributionMeasure) = MeasureBase.mspace_elsize(m.obj)
 @inline MeasureBase.mspace_flatsize(m::DistributionMeasure) = MeasureBase.mspace_flatsize(m.obj)
 @inline MeasureBase.mspace_flatsize(::Type{<:Distribution{Univariate}}) = ()
+@inline MeasureBase.mspace_ndims(::Type{<:Distribution{<:ArrayLikeVariate{N}}}) where {N} = N
+@inline MeasureBase.mspace_ndims(::Type{AsMeasure{D}}) where {D<:Distribution} = MeasureBase.mspace_ndims(D)
 @inline MeasureBase.mspace_flatsize(::Type{AsMeasure{D}}) where {D<:Distribution} = MeasureBase.mspace_flatsize(D)
 
 @inline MeasureBase.preferred_stdmeasure(::Type{AsMeasure{D}}) where {D<:Distribution} = MeasureBase.preferred_stdmeasure(D)
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 5aa2eb09..f4e87bb8 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -198,6 +198,7 @@ include("mass-interface.jl")
 
 include("density.jl")
 include("density-core.jl")
+include("density-batched.jl")
 
 include("primitives/counting.jl")
 include("primitives/lebesgue.jl")
@@ -210,7 +211,6 @@ include("combinators/weighted.jl")
 include("combinators/superpose.jl")
 include("combinators/product.jl")
 include("combinators/power.jl")
-include("density-batched.jl")
 include("transport-batched.jl")
 include("combinators/spikemixture.jl")
 include("combinators/likelihood.jl")
diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl
index 2c24ca7f..a2134709 100644
--- a/src/combinators/bind.jl
+++ b/src/combinators/bind.jl
@@ -305,10 +305,8 @@ function _bind_ld_impl(::Type{Pair}, μ::Bind, xy::Pair)
 end
 
 function _bind_ld_impl(::Union{typeof(vcat),typeof(merge)}, μ::Bind, xy)
-    ℓ, x_μ, x_rest = logdensityof_with_rest(μ, xy)
-    if !isempty(x_rest)
-        throw(ArgumentError("Variate too long during density evaluation of a bind"))
-    end
+    ℓ, _, x_rest = logdensityof_with_rest(μ, xy)
+    isempty(x_rest) || _throw_stream_too_long()
     return ℓ
 end
 
@@ -320,6 +318,11 @@ function _bind_ld_impl(@nospecialize(f_c), @nospecialize(μ::Bind), @nospecializ
     )
 end
 
+# The secondary measure depends on the primary variate, so streams are
+# consumed one by one:
+@inline fixed_stream_size(::Type{<:Bind}) = static(false)
+@inline mspace_ndims(::Type{<:_BindBy{typeof(vcat)}}) = 1
+
 function logdensityof_with_rest(μ::_BindBy{typeof(vcat)}, x::AbstractVector)
     ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x)
     ℓ_b, b, x_rest = logdensityof_with_rest(_get_β_a(μ, a), x2)
@@ -333,6 +336,14 @@ function logdensityof_with_rest(μ::_BindBy{typeof(merge)}, x::NamedTuple)
     return ℓ_a + ℓ_b, merge(a, b), x_rest
 end
 
+function batched_logdensityof_with_rest(μ::_BindBy{typeof(vcat)}, x::AbstractVector, ::Tuple{})
+    ℓ, _, x_rest = logdensityof_with_rest(μ, x)
+    return ℓ, x_rest
+end
+
+batched_logdensityof_impl(μ::_BindBy{typeof(vcat)}, X::AbstractArray) = _streamwise_ld(logdensityof_impl, μ, X)
+batched_logdensityof_impl(μ::_BindBy{typeof(vcat)}, x::AbstractVector) = _bind_ld_impl(vcat, μ, x)
+
 
 function rand_impl(ctx::GenContext, μ::Bind)
     a = rand_impl(ctx, μ.α)
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index 68cd16d6..55c92bce 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -169,15 +169,13 @@ function _combined_ld_impl(::Type{Pair}, μ::CombinedMeasure, ab::Pair)
     logdensityof(μ.α, ab.first) + logdensityof(μ.β, ab.second)
 end
 
-function _combined_ld_impl(::Union{typeof(vcat),typeof(merge)}, μ::CombinedMeasure, ab)
-    ℓ, x_μ, x_rest = logdensityof_with_rest(μ, ab)
-    if !isempty(x_rest)
-        throw(
-            ArgumentError(
-                "Variate too long during density evaluation of a combined measure",
-            ),
-        )
-    end
+function _combined_ld_impl(::typeof(vcat), μ::CombinedMeasure, ab::AbstractVector)
+    _point_result(_materialize(_combined_batched_ld(μ, ab, static(true))), μ)
+end
+
+function _combined_ld_impl(::typeof(merge), μ::CombinedMeasure, ab::NamedTuple)
+    ℓ, _, x_rest = logdensityof_with_rest(μ, ab)
+    isempty(x_rest) || _throw_stream_too_long()
     return ℓ
 end
 
@@ -186,20 +184,51 @@ function _combined_ld_impl(f_c, μ::CombinedMeasure, ab)
     return logdensityof(tpm_α, a) + logdensityof(μ.β, b)
 end
 
-function batched_logdensityof_impl(μ::CombinedMeasure{typeof(vcat)}, A::AbstractArray)
-    ℓ, _, A_rest = batched_logdensityof_with_rest(μ, A)
-    if size(A_rest, 1) != 0
-        throw(ArgumentError("Variate streams too long during batched density evaluation of a combined measure"))
-    end
+@inline mspace_ndims(::Type{<:CombinedMeasure{typeof(vcat)}}) = 1
+@inline function fixed_stream_size(::Type{<:CombinedMeasure{<:Any,MA,MB}}) where {MA,MB}
+    static(fixed_stream_size(MA) === static(true) && fixed_stream_size(MB) === static(true))
+end
+
+# Batches of vcat-combined variates are batches of streams: with fixed
+# component sizes the whole batch is consumed in fused operations,
+# otherwise stream by stream.
+@inline function batched_logdensityof_impl(μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray)
+    _combined_batched_ld(μ, X, fixed_stream_size(μ))
+end
+function _combined_batched_ld(μ::CombinedMeasure, X::AbstractArray, ::True)
+    ℓ, X_rest = batched_logdensityof_with_rest(μ, X, ())
+    size(X_rest, 1) == 0 || _throw_stream_too_long()
     return ℓ
 end
+_combined_batched_ld(μ::CombinedMeasure, X::AbstractVector, ::False) = _combined_batched_ld(μ, X, static(true))
+_combined_batched_ld(μ::CombinedMeasure, X::AbstractArray, ::False) = _streamwise_ld(logdensityof_impl, μ, X)
+
+function batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray, ::Tuple{})
+    _combined_ld_with_rest(μ, X)
+end
+batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector, ::Tuple{}) = _combined_ld_with_rest(μ, x)
+function _combined_ld_with_rest(μ::CombinedMeasure, X::AbstractArray)
+    ℓ_a, X2 = batched_logdensityof_with_rest(μ.α, X, ())
+    ℓ_b, X_rest = batched_logdensityof_with_rest(μ.β, X2, ())
+    return _lazy_add(ℓ_a, ℓ_b), X_rest
+end
+
+# Several variates per stream interleave the component parts, so the rows
+# of each variate are split by the fixed component sizes:
+function batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray, sz::Dims)
+    n_a, n_b = _fixed_stream_length(μ.α), _fixed_stream_length(μ.β)
+    X_μ, X_rest = _batched_split(X, (n_a + n_b) * prod(sz))
+    X_v = reshape(X_μ, (n_a + n_b, sz..., Base.tail(size(X_μ))...))
+    X_a, X_b = _batched_split(X_v, n_a)
+    ℓ_a, _ = batched_logdensityof_with_rest(μ.α, X_a, ())
+    ℓ_b, _ = batched_logdensityof_with_rest(μ.β, X_b, ())
+    return _lazy_add(ℓ_a, ℓ_b), X_rest
+end
 
-function batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, A::AbstractArray)
-    ℓ_a, _, A2 = batched_logdensityof_with_rest(μ.α, A)
-    ℓ_b, _, A_rest = batched_logdensityof_with_rest(μ.β, A2)
-    n_μ = size(A, 1) - size(A_rest, 1)
-    A_μ = view(A, 1:n_μ, Base.tail(axes(A))...)
-    return _lazy_add(ℓ_a, ℓ_b), A_μ, A_rest
+@inline _fixed_stream_length(μ) = _fixed_stream_length(μ, mspace_flatsize(μ))
+@inline _fixed_stream_length(μ, sz::SizeLike) = dynamic(size2length(sz))
+@noinline function _fixed_stream_length(μ, ::NoMSpaceElementSize)
+    throw(ArgumentError("Consuming several variates per stream requires measures of type $(nameof(typeof(μ))) to have a known variate size"))
 end
 
 function logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector)
diff --git a/src/combinators/half.jl b/src/combinators/half.jl
index 3873a86c..c6ed31da 100644
--- a/src/combinators/half.jl
+++ b/src/combinators/half.jl
@@ -7,6 +7,8 @@ end
 @inline mspace_elsize(μ::Half) = mspace_elsize(μ.parent)
 @inline mspace_flatsize(μ::Half) = mspace_flatsize(μ.parent)
 @inline mspace_flatsize(::Type{<:Half{M}}) where {M} = mspace_flatsize(M)
+@inline mspace_ndims(::Type{<:Half{M}}) where {M} = mspace_ndims(M)
+@inline fixed_stream_size(::Type{<:Half{M}}) where {M} = fixed_stream_size(M)
 @inline preferred_stdmeasure(::Type{<:Half}) = StdUniform
 
 function Base.show(io::IO, μ::Half)
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index b48f38f5..49f6ba00 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -122,12 +122,81 @@ params(d::PowerMeasure) = params(first(marginals(d)))
     basemeasure(d.parent)^d.axes
 end
 
-# Densities of powers are evaluated by the batched density machinery over
-# the flat variate storage (see density-batched.jl):
+@inline mspace_ndims(::Type{<:PowerMeasure{M,A}}) where {M,A<:Tuple} = _add_ndims(mspace_ndims(M), fieldcount(A))
+@inline fixed_stream_size(::Type{<:PowerMeasure{M}}) where {M} = fixed_stream_size(M)
+
+# The innermost base measure of nested powers and the total number of power
+# dimensions:
+@inline _pwr_unwrap(μ) = (μ, static(0))
+@inline function _pwr_unwrap(μ::PowerMeasure)
+    ν, n = _pwr_unwrap(pwr_base(μ))
+    ν, n + static(length(pwr_axes(μ)))
+end
+
+# Batched kernels: the base kernel runs over the flat batch, the power then
+# sums the leading dimensions of the result that belong to its axes.
+@inline function _powered_kernel(f::F, μ::PowerMeasure, X) where {F}
+    _check_pwr_batch(X, mspace_flatsize(μ))
+    _sum_leading_dims(_batched_kernel(f, pwr_base(μ), X), static(length(pwr_axes(μ))))
+end
+@inline _check_pwr_batch(X::AbstractArray, sz_flat::SizeLike) = _check_flatsize(X, sz_flat)
+@inline _check_pwr_batch(X, ::Any) = nothing
+@inline batched_logdensityof_impl(μ::PowerMeasure, X) = _powered_kernel(logdensityof_impl, μ, X)
+@inline batched_logdensity_def(μ::PowerMeasure, X) = _powered_kernel(logdensity_def, μ, X)
+
+# Point evaluation: flat variates are batches with zero batch dimensions,
+# nested variates without flat storage sum the point densities of the base.
+@inline _point_ld(f::F, μ::PowerMeasure, x::AbstractArray{<:Number}) where {F} = f(μ, x)
+@inline logdensityof_impl(μ::PowerMeasure, x) = _powered_point(logdensityof_impl, μ, x)
+@inline logdensity_def(μ::PowerMeasure, x) = _powered_point(logdensity_def, μ, x)
+
+@inline function _powered_point(f::F, μ::PowerMeasure, x::AbstractArray{<:Number}) where {F}
+    _check_pwr_variate(μ, x, mspace_flatsize(μ))
+    _point_result(_materialize(_batched_kernel(f, μ, x)), μ)
+end
+@inline function _powered_point(f::F, μ::PowerMeasure, x::AbstractArray) where {F}
+    _powered_point_nested(f, μ, x, _flat_storage(x))
+end
+@inline function _powered_point_nested(f::F, μ::PowerMeasure, x, x_flat::AbstractArray) where {F}
+    _check_pwr_variate(μ, x, mspace_flatsize(μ))
+    _point_result(_materialize(_batched_kernel(f, μ, x_flat)), μ)
+end
+function _powered_point_nested(f::F, μ::PowerMeasure, x::AbstractArray, ::NoFlatStorage) where {F}
+    if maybestatic_size(x) != pwr_size(μ)
+        _throw_size_mismatch()
+    end
+    ν = pwr_base(μ)
+    sum(_PointLogd(f, ν), x; init = zero(_logd_numtype(x)))
+end
+@noinline function _powered_point(::F, ::PowerMeasure, x) where {F}
+    throw(ArgumentError("Variates of powers of measures must be arrays"))
+end
 
-@inline logdensityof_impl(μ::PowerMeasure, x) = _powered_ld(logdensityof_impl, μ, x)
-@inline logdensity_def(μ::PowerMeasure, x) = _powered_ld(logdensity_def, μ, x)
-@inline batched_logdensityof_impl(μ::PowerMeasure, A::AbstractArray) = _batched_ld(logdensityof_impl, μ, A)
+# Flat variates must match the flat size where it is known, nested variates
+# the power's shape:
+@inline function _check_pwr_variate(μ::PowerMeasure, x::AbstractArray{<:Number}, sz_flat::SizeLike)
+    if !_matches_flatsize(maybestatic_size(x), sz_flat) && maybestatic_size(x) != pwr_size(μ)
+        _throw_size_mismatch()
+    end
+    return nothing
+end
+@inline function _check_pwr_variate(μ::PowerMeasure, x::AbstractArray, ::Any)
+    if maybestatic_size(x) != pwr_size(μ)
+        _throw_size_mismatch()
+    end
+    return nothing
+end
+
+# Streams: a power consumes its size times the variates of the base measure
+# and sums the base results over its axes.
+function batched_logdensityof_with_rest(μ::PowerMeasure, X::AbstractArray, sz::Dims)
+    _powered_ld_with_rest(μ, X, sz)
+end
+batched_logdensityof_with_rest(μ::PowerMeasure, x::AbstractVector, sz::Tuple{}) = _powered_ld_with_rest(μ, x, sz)
+function _powered_ld_with_rest(μ::PowerMeasure, X::AbstractArray, sz::Dims)
+    ℓ, X_rest = batched_logdensityof_with_rest(pwr_base(μ), X, (map(dynamic, _size_dims(pwr_size(μ)))..., sz...))
+    return _sum_leading_dims(ℓ, static(length(pwr_axes(μ)))), X_rest
+end
 
 # Support checks of powers run over the flat variate storage where the base
 # measure has scalar variates, elementwise otherwise:
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 1d46e93a..c1f06fe7 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -40,6 +40,7 @@ for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :log
     end
 end
 
+
 # Variates of products are collections of marginal variates, with the same
 # structure as the marginals:
 @inline function _check_marginal_count(mar::AbstractArray, x::AbstractArray)
@@ -66,6 +67,18 @@ end
 proxy(μ::ProductMeasure{<:FillArrays.Fill}) =
     powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ)))
 
+# Batches of tuple and named tuple variates are tuples resp. named tuples
+# of batches, the marginal densities add up lazily:
+for (bhead, head) in [(:batched_logdensityof_impl, :logdensityof_impl), (:batched_logdensity_def, :logdensity_def)]
+    @eval @inline function $bhead(μ::ProductMeasure{<:Tuple}, X::Tuple)
+        _lazy_sum(map((m, Xi) -> _batched_kernel($head, m, Xi), marginals(μ), X))
+    end
+    @eval @inline function $bhead(μ::ProductMeasure{<:NamedTuple{names}}, X::NamedTuple{names}) where {names}
+        _lazy_sum(map((m, Xi) -> _batched_kernel($head, m, Xi), values(marginals(μ)), values(X)))
+    end
+end
+@inline _lazy_sum(ℓs::Tuple) = reduce(_lazy_add, ℓs)
+
 # Relative densities between products evaluate marginal-wise. Support
 # checks happen at the logdensity_rel level for the whole products, so the
 # unsafe marginal evaluation suffices here:
@@ -187,35 +200,50 @@ marginals(μ::ProductMeasure) = μ.marginals
     _cat_sizes(mspace_flatsize(M), maybestatic_size(marginals(μ)))
 end
 
-# Batched densities over flat storage `(marginal flat dims..., product
+@inline function mspace_ndims(::Type{<:ProductMeasure{<:AbstractArray{M,N}}}) where {M,N}
+    _add_ndims(mspace_ndims(M), N)
+end
+@inline fixed_stream_size(::Type{<:ProductMeasure{<:AbstractArray{M}}}) where {M} = fixed_stream_size(M)
+
+# Batched kernels over flat storage `(marginal variate dims..., product
 # dims..., batch dims...)`. Marginals with scalar variates align with the
 # leading dimensions of the batch, so one broadcast evaluates all marginal
 # densities. Marginals with array variates are evaluated one by one over
 # their slices of the batch.
-@inline function batched_logdensityof_impl(μ::ProductMeasure{<:AbstractArray{M,N}}, A::AbstractArray) where {M,N}
-    _product_batched_ld(μ, A, mspace_flatsize(M), Val(N), Val(isconcretetype(M)))
+for (bhead, head) in [(:batched_logdensityof_impl, :logdensityof_impl), (:batched_logdensity_def, :logdensity_def)]
+    @eval @inline function $bhead(μ::ProductMeasure{<:AbstractArray{M}}, X::AbstractArray) where {M}
+        _array_product_kernel($head, μ, X, _static_ndims_of(mspace_ndims(M)))
+    end
 end
 
-@inline function _product_batched_ld(μ::ProductMeasure, A::AbstractArray, ::Tuple{}, ::Val{N}, ::Val{true}) where {N}
-    ld = Broadcast.instantiate(Broadcast.broadcasted(dynamic ∘ logdensityof_impl, marginals(μ), A))
-    _sum_leading_dims(ld, static(N))
+@inline function _array_product_kernel(f::F, μ::ProductMeasure, X::AbstractArray, k::Integer) where {F}
+    _array_product_kernel(f, μ, X, static(k))
 end
-
-@inline function _product_batched_ld(μ::ProductMeasure, A::AbstractArray, sz_m::SizeLike, ::Val{N}, ::Val{true}) where {N}
-    _marginal_slices_ld(marginals(μ), A, Val(length(sz_m)), Val(ndims(A) - length(sz_m) - N))
+@inline function _array_product_kernel(f::F, μ::ProductMeasure, X::AbstractArray, ::StaticInteger{0}) where {F}
+    mar = marginals(μ)
+    _check_flatsize(X, maybestatic_size(mar))
+    ld = Broadcast.instantiate(Broadcast.broadcasted(_DynamicPointLogd(f), mar, X))
+    _sum_leading_dims(ld, static(ndims(mar)))
+end
+@inline function _array_product_kernel(f::F, μ::ProductMeasure, X::AbstractArray, ::StaticInteger{K}) where {F,K}
+    _marginal_slices_ld(f, marginals(μ), X, Val(K), Val(ndims(X) - K - ndims(marginals(μ))))
+end
+@noinline function _array_product_kernel(::F, μ::ProductMeasure{<:AbstractArray{M}}, ::AbstractArray, ::NoMSpaceElementSize) where {F,M}
+    throw(ArgumentError("Batched density evaluation of products over arrays of marginals of type $(nameof(M)) requires MeasureBase.mspace_ndims to be declared for that type"))
 end
 
-@inline function _product_batched_ld(μ::ProductMeasure, A::AbstractArray, ::Any, ::Val, ::Val)
-    _batched_ld_generic(logdensityof_impl, μ, A)
+struct _DynamicPointLogd{F} <: Function
+    f::F
 end
+@inline (k::_DynamicPointLogd)(m, x) = _dynamic_logd(k.f(m, x), x)
 
-function _marginal_slices_ld(mar::AbstractArray{<:Any,N}, A::AbstractArray, ::Val{K}, ::Val{B}) where {N,K,B}
+function _marginal_slices_ld(f::F, mar::AbstractArray{<:Any,N}, A::AbstractArray, ::Val{K}, ::Val{B}) where {F,N,K,B}
     if ndims(A) != K + N + B || ntuple(i -> size(A, K + i), Val(N)) != size(mar)
         _throw_size_mismatch()
     end
     lead = ntuple(_ -> Colon(), Val(K))
     trail = ntuple(_ -> Colon(), Val(B))
-    ld(i) = _materialize(batched_logdensityof_impl(mar[i], view(A, lead..., Tuple(i)..., trail...)))
+    ld(i) = _materialize(_batched_kernel(f, mar[i], view(A, lead..., Tuple(i)..., trail...)))
     init = _zero_logd(A, ntuple(i -> size(A, K + N + i), Val(B)))
     return mapreduce(ld, +, CartesianIndices(mar); init = init)
 end
@@ -223,20 +251,27 @@ end
 @inline _zero_logd(A::AbstractArray, ::Tuple{}) = zero(_logd_numtype(A))
 @inline _zero_logd(A::AbstractArray, dims::Tuple) = fill!(similar(A, _logd_numtype(A), dims), 0)
 
-# The point density of array products with array-variate marginals accepts
-# the flat variate storage `(marginal flat dims..., product dims...)`:
+# Point densities of array products at numeric variates go through the
+# batched kernel where the variate rank of the marginals is known,
+# marginal by marginal otherwise:
+@inline _point_ld(f::F, μ::AbstractProductMeasure, x::AbstractArray{<:Number}) where {F} = f(μ, x)
 @inline function logdensityof_impl(μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray{<:Number}) where {M}
-    _array_product_ld(μ, x, mspace_flatsize(M))
+    _array_product_ld(logdensityof_impl, μ, x, mspace_ndims(M))
+end
+@inline function logdensity_def(μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray{<:Number}) where {M}
+    _array_product_ld(logdensity_def, μ, x, mspace_ndims(M))
 end
-@inline _array_product_ld(μ::ProductMeasure, x::AbstractArray, ::Tuple{}) = _array_product_ld_nested(μ, x)
-@inline _array_product_ld(μ::ProductMeasure, x::AbstractArray, ::NoMSpaceElementSize) = _array_product_ld_nested(μ, x)
-@inline function _array_product_ld(μ::ProductMeasure, x::AbstractArray, sz_m::SizeLike)
-    _marginal_slices_ld(marginals(μ), x, Val(length(sz_m)), Val(0))
+@inline function _array_product_ld(f::F, μ::ProductMeasure, x::AbstractArray, ::Integer) where {F}
+    _point_result(_materialize(_batched_kernel(f, μ, x)), μ)
 end
-@inline function _array_product_ld_nested(μ::ProductMeasure, x::AbstractArray)
+@inline function _array_product_ld(f::F, μ::ProductMeasure, x::AbstractArray, ::NoMSpaceElementSize) where {F}
+    _array_product_ld_nested(f, μ, x)
+end
+@inline function _array_product_ld_nested(f::F, μ::ProductMeasure, x::AbstractArray) where {F}
     _check_marginal_count(marginals(μ), x)
-    mapreduce(logdensityof, +, marginals(μ), x)
+    mapreduce(_PointLogd(f, nothing), +, marginals(μ), x)
 end
+@inline (k::_PointLogd{F,Nothing})(m, x) where {F} = _point_ld(k.f, m, x)
 
 # TODO: Better `map` support in MappedArrays
 _map(f, args...) = map(f, args...)
@@ -445,3 +480,34 @@ end
 function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{false}) where {S}
     _batched_from_std(S, μ, Z, mspace_flatsize(μ))
 end
+
+
+# Streams: tuple products consume marginal by marginal, so marginals of
+# value-dependent size are supported for a single variate per stream.
+function batched_logdensityof_with_rest(μ::ProductMeasure{<:Tuple}, X::AbstractArray, ::Tuple{})
+    _marginals_ld_with_rest(marginals(μ), X)
+end
+function batched_logdensityof_with_rest(μ::ProductMeasure{<:Tuple}, x::AbstractVector, ::Tuple{})
+    _marginals_ld_with_rest(marginals(μ), x)
+end
+function _marginals_ld_with_rest(ms::Tuple, X::AbstractArray)
+    ℓ1, X2 = batched_logdensityof_with_rest(ms[1], X, ())
+    ℓ_rest, X_rest = _marginals_ld_with_rest(Base.tail(ms), X2)
+    return _lazy_add(ℓ1, ℓ_rest), X_rest
+end
+function _marginals_ld_with_rest(ms::Tuple{Any}, X::AbstractArray)
+    batched_logdensityof_with_rest(ms[1], X, ())
+end
+function batched_logdensityof_with_rest(μ::ProductMeasure{<:NamedTuple{names}}, X::AbstractArray, sz::Tuple{}) where {names}
+    batched_logdensityof_with_rest(productmeasure(values(marginals(μ))), X, sz)
+end
+function batched_logdensityof_with_rest(μ::ProductMeasure{<:NamedTuple{names}}, x::AbstractVector, sz::Tuple{}) where {names}
+    batched_logdensityof_with_rest(productmeasure(values(marginals(μ))), x, sz)
+end
+
+@inline function fixed_stream_size(::Type{<:ProductMeasure{M}}) where {M<:Tuple}
+    static(all(T -> fixed_stream_size(T) === static(true), M.parameters))
+end
+@inline function fixed_stream_size(::Type{<:ProductMeasure{NamedTuple{names,M}}}) where {names,M<:Tuple}
+    fixed_stream_size(ProductMeasure{M})
+end
diff --git a/src/combinators/restricted.jl b/src/combinators/restricted.jl
index f80e4dc4..b4339a59 100644
--- a/src/combinators/restricted.jl
+++ b/src/combinators/restricted.jl
@@ -6,6 +6,8 @@ end
 @inline mspace_elsize(μ::RestrictedMeasure) = mspace_elsize(μ.base)
 @inline mspace_flatsize(μ::RestrictedMeasure) = mspace_flatsize(μ.base)
 @inline mspace_flatsize(::Type{<:RestrictedMeasure{<:Any,M}}) where {M} = mspace_flatsize(M)
+@inline mspace_ndims(::Type{<:RestrictedMeasure{<:Any,M}}) where {M} = mspace_ndims(M)
+@inline fixed_stream_size(::Type{<:RestrictedMeasure{<:Any,M}}) where {M} = fixed_stream_size(M)
 
 @inline logdensity_def(d::RestrictedMeasure, x) = logdensity_def(d.base, x)
 
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index 3ca5d554..2da33ce1 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -172,6 +172,17 @@ end
 
 
 @inline mspace_flatsize(μ::SuperpositionMeasure) = mspace_flatsize(typeof(μ))
+
+# The variate rank of a superposition is the common rank of its components:
+@inline mspace_ndims(::Type{<:SuperpositionMeasure{C}}) where {C<:AbstractArray} = mspace_ndims(eltype(C))
+@generated function mspace_ndims(::Type{MU}) where {C<:Tuple,MU<:SuperpositionMeasure{C}}
+    args = [:(mspace_ndims($T)) for T in C.parameters]
+    :(_common_ndims(($(args...),), MU))
+end
+@inline function _common_ndims(ns::Tuple{Integer,Vararg{Integer}}, ::Type{MU}) where {MU}
+    all(==(first(ns)), ns) ? first(ns) : NoMSpaceElementSize{MU}()
+end
+@inline _common_ndims(::Tuple, ::Type{MU}) where {MU} = NoMSpaceElementSize{MU}()
 @inline mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:AbstractArray} = _scalar_or_unknown(mspace_flatsize(eltype(C)))
 @inline mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:Tuple} = _common_scalar_flatsize(C)
 @generated function _common_scalar_flatsize(::Type{C}) where {C<:Tuple}
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index b8d5d0c4..9e9682de 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -72,25 +72,37 @@ export PushforwardMeasure
     Users should not call `PushforwardMeasure` directly. Instead call or add
     methods to `pushfwd`.
 """
-struct PushforwardMeasure{F,I,M,S<:PushFwdStyle} <: AbstractPushforward
+struct PushforwardMeasure{F,I,M,S<:PushFwdStyle,VS} <: AbstractPushforward
     f::F
     finv::I
     origin::M
     style::S
+    varsize::VS
 
-    function PushforwardMeasure{F,I,M,S}(
-        f::F,
-        finv::I,
-        origin::M,
-        style::S,
-    ) where {F,I,M,S<:PushFwdStyle}
-        new{F,I,M,S}(f, finv, origin, style)
+    function PushforwardMeasure(f, finv, origin::M, style::S, varsize::VS) where {M,S<:PushFwdStyle,VS}
+        new{Core.Typeof(f),Core.Typeof(finv),M,S,VS}(f, finv, origin, style, varsize)
     end
+end
 
-    function PushforwardMeasure(f, finv, origin::M, style::S) where {M,S<:PushFwdStyle}
-        new{Core.Typeof(f),Core.Typeof(finv),M,S}(f, finv, origin, style)
-    end
+# The size of the variates of a pushforward follows from a test value of
+# the origin, where the origin has variates of known size:
+@inline function _pushfwd_varsize(f, μ)
+    _pushfwd_varsize(f, μ, mspace_flatsize(μ))
+end
+@inline _pushfwd_varsize(f, μ, ::SizeLike) = _value_flatsize(f(testvalue(μ)))
+@inline _pushfwd_varsize(f, μ, ::NoMSpaceElementSize) = NoMSpaceElementSize{typeof(μ)}()
+
+@inline mspace_elsize(ν::PushforwardMeasure) = _value_or_unknown(ν.varsize, ν)
+@inline mspace_flatsize(ν::PushforwardMeasure) = _value_or_unknown(ν.varsize, ν)
+@inline _value_or_unknown(sz::SizeLike, ν) = sz
+@inline _value_or_unknown(::NoMSpaceElementSize, ν) = NoMSpaceElementSize{typeof(ν)}()
+@inline fixed_stream_size(::Type{<:PushforwardMeasure{<:Any,<:Any,<:Any,<:Any,VS}}) where {VS} = static(VS <: SizeLike)
+@inline function mspace_ndims(::Type{MU}) where {VS,MU<:PushforwardMeasure{<:Any,<:Any,<:Any,<:Any,VS}}
+    _ndims_of_size_type(VS, MU)
 end
+@inline _ndims_of_size_type(::Type{<:Tuple{Vararg{Any,N}}}, ::Type) where {N} = N
+@inline _ndims_of_size_type(::Type{StaticArrays.Size{S}}, ::Type) where {S} = length(S)
+@inline _ndims_of_size_type(::Type, ::Type{MU}) where {MU} = NoMSpaceElementSize{MU}()
 
 const _NonBijectivePusfwdMeasure{M<:PushforwardMeasure,S<:PushFwdStyle} = Union{
     PushforwardMeasure{<:Any,<:NoInverse,M,S},
@@ -272,7 +284,7 @@ export pushfwd
 @inline pushfwd(::typeof(identity), μ) = μ
 @inline pushfwd(::typeof(identity), μ, ::PushFwdStyle) = μ
 
-_pushfwd_impl(f, μ, style) = PushforwardMeasure(f, inverse(f), μ, style)
+_pushfwd_impl(f, μ, style) = PushforwardMeasure(f, inverse(f), μ, style, _pushfwd_varsize(f, μ))
 
 function _pushfwd_impl(
     f,
@@ -282,7 +294,7 @@ function _pushfwd_impl(
     orig_μ = μ.origin
     new_f = fcomp(f, μ.f)
     new_f_inv = fcomp(μ.finv, inverse(f))
-    PushforwardMeasure(new_f, new_f_inv, orig_μ, style)
+    PushforwardMeasure(new_f, new_f_inv, orig_μ, style, _pushfwd_varsize(new_f, orig_μ))
 end
 
 # Simplifications for Dirac and WeightedMeasure origins are defined in
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index b9a8a89b..fba7806a 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -26,8 +26,11 @@ _logweight(::AbstractMeasure) = 0
     _logweight_for(d.logweight, x) + logdensityof_impl(basemeasure(d), x)
 end
 
-@inline function batched_logdensityof_impl(d::AbstractWeightedMeasure, A::AbstractArray)
-    _lazy_add(_logweight_for(d.logweight, A), batched_logdensityof_impl(basemeasure(d), A))
+@inline function batched_logdensityof_impl(d::AbstractWeightedMeasure, X)
+    _lazy_add(_logweight_for(d.logweight, X), batched_logdensityof_impl(basemeasure(d), X))
+end
+@inline function batched_logdensity_def(d::AbstractWeightedMeasure, X)
+    _lazy_add(_logweight_for(d.logweight, X), _zero_logd_batch(X, mspace_ndims(basemeasure(d))))
 end
 
 @inline rand_impl(ctx::GenContext, μ::AbstractWeightedMeasure) = rand_impl(ctx, basemeasure(μ))
@@ -46,6 +49,8 @@ end
 @inline mspace_elsize(μ::WeightedMeasure) = mspace_elsize(μ.base)
 @inline mspace_flatsize(μ::WeightedMeasure) = mspace_flatsize(μ.base)
 @inline mspace_flatsize(::Type{<:WeightedMeasure{<:Any,M}}) where {M} = mspace_flatsize(M)
+@inline mspace_ndims(::Type{<:WeightedMeasure{<:Any,M}}) where {M} = mspace_ndims(M)
+@inline fixed_stream_size(::Type{<:WeightedMeasure{<:Any,M}}) where {M} = fixed_stream_size(M)
 
 massof(w::WeightedMeasure) = exp(w.logweight) * massof(w.base)
 
diff --git a/src/density-batched.jl b/src/density-batched.jl
index 9ad63765..e9c75499 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -1,236 +1,208 @@
 # This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
 
+# Batched-first density evaluation.
+#
+# A flat batch of variates is an array `(variate dims..., batch dims...)`,
+# zero batch dims meaning a single variate. Nested batches (arrays of
+# variates with flat storage, see ArraysOfArrays) are fused at the entry
+# points, the container then determines the batch dimensions. Batched
+# kernels know the variate rank of their measure (see `mspace_ndims`) and
+# return arrays over the batch dimensions, a number for a single variate.
+# Structural measures implement their kernels once, in terms of the kernels
+# of their components.
+
 export logdensities
 
 """
-    logdensities(μ::AbstractMeasure, X::AbstractArray)
-
-Compute the log-density of `μ` at each point in `X`.
+    logdensities(μ::AbstractMeasure, X)
 
-Returns an array of the shape of `X`, semantically equivalent to
-`logdensityof.(Ref(μ), X)`. Batches with flat storage of a known variate
-size (e.g. `ArraysOfArrays.ArrayOfSimilarArrays`) are evaluated in one
-fused operation over the flat data, compatible with GPU-backed storage.
+Compute the log-density of `μ` at each variate in the batch `X`.
 
-For measures with array-valued variates, `X` may also be the flat storage
-of the batch itself, with the variate dimensions leading (see
-[`MeasureBase.mspace_flatsize`](@ref)). The result then has the remaining
-dimensions of `X`.
+`X` is an array of variates (e.g. an `ArraysOfArrays.ArrayOfSimilarArrays`,
+or an array of numbers for measures with scalar variates), the flat storage
+of a batch (an array of numbers with the variate dimensions leading, see
+[`MeasureBase.mspace_ndims`](@ref)), or a tuple resp. `NamedTuple` of
+batches for measures with tuple resp. `NamedTuple` variates. Returns an
+array over the batch dimensions, semantically equivalent to
+`logdensityof.(Ref(μ), X)` for arrays of variates. Batches with flat
+storage are evaluated in fused operations, compatible with GPU and traced
+arrays.
 
-Measure types should specialize
-[`MeasureBase.batched_logdensityof_impl`](@ref) instead of `logdensities`
-itself.
+Measure types implement [`MeasureBase.batched_logdensityof_impl`](@ref).
 """
 function logdensities end
 
-@inline logdensities(μ::AbstractMeasure, X::AbstractArray) = _materialize(_batched_ld(logdensityof_impl, μ, X))
+@inline logdensities(μ::AbstractMeasure, X) = _materialize(_batched_ld(logdensityof_impl, μ, X))
 
 """
-    MeasureBase.batched_logdensityof_impl(μ::AbstractMeasure, A::AbstractArray)
-
-Implements [`logdensities`](@ref) for a batch `A` of variates of `μ` in
-flat storage: the leading dimensions of `A` are the variate dimensions
-(see [`MeasureBase.mspace_flatsize`](@ref)), any further dimensions are
-batch dimensions. Returns the log-densities as an array over the batch
-dimensions, or a scalar if there are none.
-
-Implementations must handle points outside the support of `μ` (the
-result must be `-Inf` there). Implementations for structural measures
-evaluate their component measures via `batched_logdensityof_impl` as well,
-power measures route back into the batched core (which unwraps their
-power structure and sums over the power dimensions).
-
-The default implementation broadcasts the log-density over `A` for
-measures with scalar variates and maps it over the variate slices of `A`
-otherwise. The result may be a lazy broadcast, callers materialize it
-where necessary.
+    MeasureBase.batched_logdensityof_impl(μ::AbstractMeasure, X)
+
+Log-densities of `μ` (relative to its root measure) at the variates of the
+flat batch `X`, an array `(variate dims..., batch dims...)`, returned as an
+array over the batch dimensions. `X` may be a single variate (zero batch
+dimensions), the result is a number then. Results may be lazy broadcasts,
+callers materialize them where necessary. The results must be `-Inf` for
+variates outside the support of `μ`.
+
+This is the primary density extension point. The default implementation
+broadcasts the point kernel [`MeasureBase.logdensityof_impl`](@ref) over
+`X` for measures with scalar variates and maps it over the variate slices
+of `X` (in a host loop) for measures with array variates of a declared
+number of dimensions (see [`MeasureBase.mspace_ndims`](@ref)). Measure
+types with array variates should implement `batched_logdensityof_impl`
+directly. Structural measures evaluate their components via
+`batched_logdensityof_impl` as well.
 """
 function batched_logdensityof_impl end
 
-@inline function batched_logdensityof_impl(μ::AbstractMeasure, A::AbstractArray)
-    _batched_ld_generic(logdensityof_impl, μ, A)
+@inline function batched_logdensityof_impl(μ::AbstractMeasure, X)
+    _default_batched_kernel(logdensityof_impl, μ, X, _static_ndims(μ))
 end
 
-# Batched density machinery, parameterized over the point-level density
-# function `f` (`logdensityof_impl` or `logdensity_def`).
-#
-# Variates and batches with flat storage are evaluated in one call of the
-# batched point kernel of the base measure: the leading dimensions of the
-# flat data are the variate dimensions of the base measure, followed by the
-# power dimensions and any batch dimensions. The power dimensions are summed
-# afterwards. Nested arrays without flat storage of a known layout are
-# evaluated level by level.
+# The variate rank as a static integer, from the type where known:
+@inline _static_ndims(μ::MU) where {MU} = _static_ndims(mspace_ndims(MU), μ)
+@inline _static_ndims(n::Integer, μ) = static(n)
+@inline _static_ndims(::NoMSpaceElementSize, μ) = _static_ndims_of(mspace_ndims(μ))
+@inline _static_ndims_of(n::Integer) = static(n)
+@inline _static_ndims_of(n::NoMSpaceElementSize) = n
 
-struct NoFlatStorage end
+"""
+    MeasureBase.batched_logdensity_def(μ::AbstractMeasure, X)
 
-@inline _batched_ld(f::F, μ, X::AbstractArray) where {F} = _batched_ld_sized(f, μ, X, mspace_flatsize(μ))
+Batched form of [`logdensity_def`](@ref): log-densities of `μ` relative to
+`basemeasure(μ)` at the variates of the flat batch `X`, with the same
+conventions and defaults as [`MeasureBase.batched_logdensityof_impl`](@ref).
+"""
+function batched_logdensity_def end
 
-@inline function _batched_ld_sized(f::F, μ, X::AbstractArray, sz_flat::SizeLike) where {F}
-    _batched_ld_flat(f, μ, X, _flat_storage(X), sz_flat)
+@inline function batched_logdensity_def(μ::AbstractMeasure, X)
+    _default_batched_kernel(logdensity_def, μ, X, _static_ndims(μ))
 end
 
-# Batches of scalar variates are arrays of numbers, nested arrays are not
-# reinterpreted as flat storage:
-@inline function _batched_ld_sized(f::F, μ, X::AbstractArray, sz_flat::Tuple{}) where {F}
-    _batched_ld_flat(f, μ, X, _flat_scalar_storage(X), sz_flat)
+@inline _default_batched_kernel(f::F, μ, X, n::Integer) where {F} = _default_batched_kernel(f, μ, X, static(n))
+@inline _default_batched_kernel(f::F, μ, X, ::StaticInteger{0}) where {F} = _scalar_kernel_broadcast(f, μ, X)
+@inline _default_batched_kernel(f::F, μ, X::AbstractArray, ::StaticInteger{0}) where {F} = _scalar_kernel_broadcast(f, μ, X)
+@inline function _default_batched_kernel(f::F, μ, X::AbstractArray, ::StaticInteger{K}) where {F,K}
+    _map_variate_slices(f, μ, X, Val(K))
 end
-
-@inline _flat_scalar_storage(X::AbstractArray{<:Number}) = X
-function _flat_scalar_storage(::AbstractArray)
-    throw(ArgumentError("A batch of scalar variates must be an array of numbers"))
+@noinline function _default_batched_kernel(f::F, μ, X, ::NoMSpaceElementSize) where {F}
+    throw(ArgumentError("Batched density evaluation for measures of type $(nameof(typeof(μ))) requires MeasureBase.mspace_ndims to be declared for the type or MeasureBase.batched_logdensityof_impl to be implemented"))
 end
 
-# Measures of unknown variate size take `X` as an array of points:
-@inline function _batched_ld_sized(f::F, μ, X::AbstractArray, ::NoMSpaceElementSize) where {F}
-    _batched_kernel(f, μ, X)
-end
-@inline function _batched_ld_sized(f::F, μ::PowerMeasure, X::AbstractArray, ::NoMSpaceElementSize) where {F}
-    Broadcast.instantiate(Broadcast.broadcasted(_DynamicLogd(_pointwise_ld_fixed(f), μ), X))
+# Point kernels of scalar-variate measures broadcast over the batch. Static
+# results are made dynamic to keep reductions type stable.
+@inline function _scalar_kernel_broadcast(f::F, μ, X::AbstractArray) where {F}
+    Broadcast.instantiate(Broadcast.broadcasted(_DynamicLogd(f, μ), X))
 end
-@inline _pointwise_ld_fixed(f::F) where {F} = (μ, x) -> _pointwise_ld(f, μ, x)
+@inline _scalar_kernel_broadcast(f::F, μ, x) where {F} = _dynamic_logd(f(μ, x), x)
 
-@inline function _batched_ld_flat(f::F, μ, X, X_flat::AbstractArray, sz_flat) where {F}
-    ν, n_pwr = _pwr_unwrap(μ)
-    _check_flatsize(X_flat, sz_flat)
-    _sum_leading_dims(_batched_kernel(f, ν, X_flat), n_pwr)
+struct _DynamicLogd{F,M} <: Function
+    f::F
+    μ::M
 end
+@inline (k::_DynamicLogd)(x) = _dynamic_logd(k.f(k.μ, x), x)
 
-@inline function _batched_ld_flat(f::F, μ, X, ::NoFlatStorage, sz_flat) where {F}
-    map(x -> _pointwise_ld(f, μ, x), X)
+# Point kernels of array-variate measures map over the variate slices of
+# the batch, a batch of a single variate is evaluated directly:
+@inline _map_variate_slices(f::F, μ, X::AbstractArray{<:Any,K}, ::Val{K}) where {F,K} = _dynamic_logd(f(μ, X), X)
+@inline function _map_variate_slices(f::F, μ, X::AbstractArray, ::Val{K}) where {F,K}
+    map(_DynamicLogd(f, μ), sliced(X, Val(K)))
 end
 
-@inline _pointwise_ld(f::F, μ, x) where {F} = f(μ, x)
-@inline _pointwise_ld(f::F, μ::PowerMeasure, x) where {F} = _powered_ld(f, μ, x)
-
-# Log-density of a power measure at a single variate:
+# The batched kernel for a point-level density function:
+@inline _batched_kernel(::typeof(logdensityof_impl), μ, X) = batched_logdensityof_impl(μ, X)
+@inline _batched_kernel(::typeof(logdensity_def), μ, X) = batched_logdensity_def(μ, X)
 
-@inline _powered_ld(f::F, μ::PowerMeasure, x) where {F} = _powered_ld_sized(f, μ, x, mspace_flatsize(μ))
-
-@inline function _powered_ld_sized(f::F, μ::PowerMeasure, x, sz_flat::SizeLike) where {F}
-    _powered_ld_flat(f, μ, x, _flat_storage(x), sz_flat)
+# Flat storage of a (nested) batch: the underlying array of memory-ordered
+# split arrays, a stacked copy for other known split modes.
+struct NoFlatStorage end
+@inline _flat_storage(X::AbstractArray{<:Number}) = X
+@inline _flat_storage(X::AbstractArray) = _flat_storage_bymode(X, getsplitmode(X))
+@inline _flat_storage(x) = NoFlatStorage()
+@inline function _flat_storage_bymode(X::AbstractArray, smode::AbstractSplitMode)
+    _flat_storage(is_memordered_splitmode(smode) ? fused(X) : stacked(X))
 end
+@inline _flat_storage_bymode(::AbstractArray, ::UnknownSplitMode) = NoFlatStorage()
+@inline _flat_storage_bymode(::AbstractArray, ::NonSplitMode) = NoFlatStorage()
 
-@inline function _powered_ld_sized(f::F, μ::PowerMeasure, x, ::NoMSpaceElementSize) where {F}
-    _powered_ld_pointwise(f, μ, x)
+# Entry: arrays of numbers are flat storage, arrays of variates are fused
+# into their flat storage (else evaluated variate by variate), tuples and
+# named tuples of batches go to the kernels directly.
+@inline _batched_ld(f::F, μ, X::AbstractArray{<:Number}) where {F} = _batched_kernel(f, μ, X)
+@inline _batched_ld(f::F, μ, X::Union{Tuple,NamedTuple}) where {F} = _batched_kernel(f, μ, X)
+@inline _batched_ld(f::F, μ, X::AbstractArray) where {F} = _batched_ld_nested(f, μ, X, _flat_storage(X), _static_ndims(μ))
+@inline function _batched_ld_nested(f::F, μ, X::AbstractArray, X_flat::AbstractArray, ::StaticInteger) where {F}
+    _check_batch_shape(_batched_kernel(f, μ, X_flat), X)
 end
-
-@inline function _powered_ld_flat(f::F, μ::PowerMeasure, x, x_flat::AbstractArray, sz_flat) where {F}
-    ν, n_pwr = _pwr_unwrap(μ)
-    if ndims(x_flat) != length(sz_flat)
-        _throw_size_mismatch()
-    end
-    _check_flatsize(x_flat, sz_flat)
-    _sum_leading_dims(_batched_kernel(f, ν, x_flat), n_pwr)
+@inline function _batched_ld_nested(f::F, μ, X::AbstractArray, ::Any, ::Any) where {F}
+    Broadcast.instantiate(Broadcast.broadcasted(_PointLogd(f, μ), X))
 end
 
-@inline _powered_ld_flat(f::F, μ::PowerMeasure, x, ::NoFlatStorage, sz_flat) where {F} = _powered_ld_pointwise(f, μ, x)
+struct _PointLogd{F,M} <: Function
+    f::F
+    μ::M
+end
+@inline (k::_PointLogd)(x) = _point_ld(k.f, k.μ, x)
 
-# Sum of the point-level densities of the base measure over the elements
-# of the variate, evaluated via the batched kernel of the base measure:
-@inline function _powered_ld_pointwise(f::F, μ::PowerMeasure, x::AbstractArray) where {F}
-    if maybestatic_size(x) != pwr_size(μ)
-        _throw_size_mismatch()
+# Results over an array of variates must have the shape of the array, a
+# mismatch means the variate rank of the measure doesn't match the batch:
+@inline function _check_batch_shape(result, X::AbstractArray)
+    if size(result) != size(X)
+        _throw_batch_shape(size(result), size(X))
     end
-    ν = pwr_base(μ)
-    _sum_leading_dims(_batched_ld_sized(f, ν, x, mspace_flatsize(ν)), static(ndims(x)))
+    return result
 end
-
-@noinline _throw_size_mismatch() = throw(ArgumentError("Size of variate doesn't match size of measure"))
-
-function _powered_ld_pointwise(f::F, μ::PowerMeasure, x) where {F}
-    throw(ArgumentError("Variates of powers of measures must be arrays, and flat variate storage requires a base measure of known variate size"))
+@noinline function _throw_batch_shape(sz_result, sz_batch)
+    throw(ArgumentError("Batched density kernel returned a result of size $sz_result for a batch of size $sz_batch, the variate dimensions of the measure don't match the batch"))
 end
 
-@inline _pointwise_ld_dyn((f, μ), x) = _dynamic_logd(_pointwise_ld(f, μ, x), x)
+# Point evaluation: array variates go through the batched kernel with zero
+# batch dimensions, other variates through the point kernel. A batched
+# kernel that returns an array for a single variate has taken the variate
+# for a batch: the variate doesn't fit the measure, or the measure lacks a
+# batched kernel for array variates.
+@inline _point_ld(f::F, μ, x::AbstractArray{<:Number}) where {F} = _point_result(_materialize(_batched_kernel(f, μ, x)), μ)
+@inline _point_ld(f::F, μ, x) where {F} = f(μ, x)
+@inline _point_ld(f::F, μ::PrimitiveMeasure, x::AbstractArray{<:Number}) where {F} = f(μ, x)
 
-@inline _pwr_unwrap(μ) = (μ, static(0))
-@inline function _pwr_unwrap(μ::PowerMeasure)
-    ν, n = _pwr_unwrap(pwr_base(μ))
-    ν, n + static(length(pwr_axes(μ)))
+@inline _point_result(ℓ::Number, μ) = ℓ
+@inline _point_result(ℓ::AbstractArray{<:Number,0}, μ) = ℓ[]
+@noinline function _point_result(ℓ, μ)
+    throw(ArgumentError("Density evaluation of measures of type $(nameof(typeof(μ))) at an array variate resulted in a batch of densities: the variate doesn't fit the measure, or the measure lacks a batched kernel for array variates"))
 end
 
-# Flat storage of a (nested) variate or batch: the underlying array of
-# memory-ordered split arrays, a stacked copy for other known split modes.
-# Nested arrays of unknown layout have no flat storage.
-
-@inline _flat_storage(x::AbstractArray{<:Number}) = x
-@inline _flat_storage(x::AbstractArray) = _flat_storage_bymode(x, getsplitmode(x))
-@inline _flat_storage(x) = NoFlatStorage()
+const _LazyBroadcast = Broadcast.Broadcasted
 
-@inline function _flat_storage_bymode(x::AbstractArray, smode::AbstractSplitMode)
-    _flat_storage(is_memordered_splitmode(smode) ? fused(x) : stacked(x))
-end
-@inline _flat_storage_bymode(::AbstractArray, ::UnknownSplitMode) = NoFlatStorage()
-@inline _flat_storage_bymode(::AbstractArray, ::NonSplitMode) = NoFlatStorage()
+@noinline _throw_size_mismatch() = throw(ArgumentError("Size of variate doesn't match size of measure"))
 
+# The leading dimensions of a flat batch must match a flat variate size:
 @inline function _check_flatsize(A::AbstractArray, sz_flat::SizeLike)
-    n = length(sz_flat)
-    if ndims(A) < n || ntuple(i -> size(A, i), Val(n)) != Tuple(sz_flat)
+    n = length(_size_dims(sz_flat))
+    if ndims(A) < n || ntuple(i -> size(A, i), Val(n)) != Tuple(_size_dims(sz_flat))
         _throw_size_mismatch()
     end
     return nothing
 end
 
-# Point kernel of the base measure over the flat batch:
-
-@inline _batched_kernel(::typeof(logdensityof_impl), ν, A::AbstractArray) = batched_logdensityof_impl(ν, A)
-@inline _batched_kernel(f::F, ν, A::AbstractArray) where {F} = _batched_ld_generic(f, ν, A)
-
-# Powers of primitive measures have log-density zero relative to their base:
-@inline function _batched_kernel(::typeof(logdensity_def), ν::PrimitiveMeasure, A::AbstractArray)
-    FillArrays.Zeros{_logd_numtype(A)}(size(A))
-end
-
-@inline _batched_ld_generic(f::F, ν, A::AbstractArray) where {F} = _batched_ld_byflatsize(f, ν, A, mspace_flatsize(ν))
-
-# Scalar variates: one lazy broadcast over the whole batch, so that
-# reductions over it don't need to allocate the intermediate result.
-# Static results of point kernels are made dynamic, to keep reductions
-# over them type stable.
-@inline function _batched_ld_byflatsize(f::F, ν, A::AbstractArray, ::Tuple{}) where {F}
-    Broadcast.instantiate(Broadcast.broadcasted(_DynamicLogd(f, ν), A))
-end
-
-struct _DynamicLogd{F,M} <: Function
-    f::F
-    ν::M
-end
-@inline (k::_DynamicLogd)(x) = _dynamic_logd(k.f(k.ν, x), x)
-
-# Array variates: map over the variate slices, or evaluate directly if `A`
-# is a single variate.
-@inline function _batched_ld_byflatsize(f::F, ν, A::AbstractArray, sz::SizeLike) where {F}
-    _batched_ld_slices(f, ν, A, Val(length(sz)))
-end
-
-# Variates of unknown size: the elements of `A` are the variates.
-@inline function _batched_ld_byflatsize(f::F, ν, A::AbstractArray, ::NoMSpaceElementSize) where {F}
-    Broadcast.instantiate(Broadcast.broadcasted(_DynamicLogd(f, ν), A))
-end
-
-@inline _batched_ld_slices(f::F, ν, A::AbstractArray{<:Any,N}, ::Val{N}) where {F,N} = f(ν, A)
-@inline function _batched_ld_slices(f::F, ν, A::AbstractArray, ::Val{M}) where {F,M}
-    map(Base.Fix1(f, ν), sliced(A, Val(M)))
-end
+@inline _materialize(bc::_LazyBroadcast) = copy(bc)
+@inline _materialize(x) = x
 
-# Sum over the leading `N` dimensions; a full reduction yields a scalar.
-# Lazy broadcasts are reduced without materialization where the broadcast
-# style supports it, and materialized before partial reductions.
 
-const _LazyBroadcast = Broadcast.Broadcasted
+# Lazy sums over the leading `N` dimensions; a full reduction yields a
+# number. Lazy broadcasts are reduced without materialization where the
+# broadcast style supports it, and materialized before partial reductions.
 const _EagerReducibleBroadcast = Broadcast.Broadcasted{<:Union{Broadcast.DefaultArrayStyle,StaticArrays.StaticArrayStyle}}
 
-@inline _materialize(bc::_LazyBroadcast) = copy(bc)
-@inline _materialize(x) = x
-
 @inline _sum_leading_dims(x::Number, ::StaticInteger{0}) = x
+@noinline function _sum_leading_dims(::Number, ::StaticInteger)
+    throw(ArgumentError("Variates of powers of measures must be arrays"))
+end
 @inline _sum_leading_dims(A::AbstractArray, n::StaticInteger) = _sum_leading_dims_impl(A, n, static(ndims(A)))
 @inline _sum_leading_dims(bc::_LazyBroadcast, n::StaticInteger) = _sum_leading_dims_lazy(bc, n, static(ndims(bc)))
 @inline _sum_leading_dims_lazy(bc::_LazyBroadcast, ::StaticInteger{0}, ::StaticInteger) = bc
 @inline _sum_leading_dims_lazy(bc::_LazyBroadcast, ::StaticInteger{0}, ::StaticInteger{0}) = bc
 @inline _sum_leading_dims_lazy(bc::_EagerReducibleBroadcast, ::StaticInteger{0}, ::StaticInteger{0}) = bc
 @inline function _sum_leading_dims_lazy(bc::_EagerReducibleBroadcast, ::StaticInteger{N}, ::StaticInteger{N}) where {N}
-    # Empty broadcasts have no known element type to reduce over lazily:
     isempty(bc) ? sum(copy(bc)) : sum(bc)
 end
 @inline _sum_leading_dims_lazy(bc::_LazyBroadcast, ::StaticInteger{N}, ::StaticInteger{N}) where {N} = sum(copy(bc))
@@ -248,63 +220,114 @@ end
     _sum_dims_seq(sum(A; dims = N), static(N - 1))
 end
 
-
 @inline _lazy_add(a::Number, b::Number) = a + b
 @inline _lazy_add(a, b) = Broadcast.instantiate(Broadcast.broadcasted(+, a, b))
 
+# Zero log-densities over the batch dimensions of a flat batch of variates
+# with `n` variate dimensions:
+@inline function _zero_logd_batch(X::AbstractArray, n::Integer)
+    FillArrays.Zeros{_logd_numtype(X)}(ntuple(i -> size(X, n + i), ndims(X) - n))
+end
+@inline _zero_logd_batch(X::AbstractArray{<:Any,N}, ::StaticInteger{N}) where {N} = zero(_logd_numtype(X))
+@inline _zero_logd_batch(x::Number, ::StaticInteger{0}) = zero(_logd_numtype(x))
 
-"""
-    MeasureBase.batched_logdensityof_with_rest(μ::AbstractMeasure, A::AbstractArray)
-
-Batched form of [`MeasureBase.logdensityof_with_rest`](@ref) for a batch
-`A` of flat vector streams: the first dimension of `A` runs along the
-streams, all further dimensions are batch dimensions.
 
-Returns a tuple `(ℓ, A_μ, A_rest)` of the log-densities over the batch
-dimensions (possibly as a lazy broadcast), the rows consumed from the
-streams and the unconsumed rest of the streams.
+# Streams: variates of composed measures are consumed from flat vector
+# streams, batched as `(rows, batch dims...)`.
 
-Consuming from streams requires the flat variate sizes of `μ` or its
-components to be known, see [`MeasureBase.mspace_flatsize`](@ref).
+"""
+    MeasureBase.batched_logdensityof_with_rest(μ::AbstractMeasure, X, sz::Dims)
+
+Consume variates of `μ` from the batch `X` of flat vector streams (first
+dimension along the streams, further dimensions are batch dimensions), a
+batch of variates of size `sz` per stream, and compute their
+log-densities.
+
+Returns a tuple `(ℓ, X_rest)` of the log-densities, an array over
+`(sz..., batch dims...)` (possibly lazy, a number for a single stream and
+`sz == ()`), and the unconsumed rest of the streams. Measure types whose
+variates have a fixed size consume `prod(sz)` variates in one batched
+kernel evaluation, the default implementation does so for the variate size
+given by [`MeasureBase.mspace_flatsize`](@ref) or
+[`MeasureBase.some_mspace_elsize`](@ref). Measure types with variates of
+value-dependent size implement `batched_logdensityof_with_rest` for
+`sz == ()` themselves, they can only be evaluated stream by stream.
 """
 function batched_logdensityof_with_rest end
 
-function batched_logdensityof_with_rest(μ::AbstractMeasure, A::AbstractArray)
-    A_μ, A_rest = _batched_consume(A, mspace_flatsize(μ))
-    return _batched_ld(logdensityof_impl, μ, A_μ), A_μ, A_rest
+function batched_logdensityof_with_rest(μ::AbstractMeasure, X::AbstractArray, sz::Dims)
+    _stream_ld_with_rest(logdensityof_impl, μ, X, sz)
 end
 
-# Consume the leading rows of a batch of streams as a batch of flat variates:
-@inline function _batched_consume(A::AbstractArray, sz::SizeLike)
-    A_flat, A_rest = _batched_split(A, dynamic(size2length(sz)))
-    return _batched_chunk_shape(A_flat, _size_dims(sz)), A_rest
+# A single stream consumes one variate via the point path:
+function batched_logdensityof_with_rest(μ::AbstractMeasure, x::AbstractVector, ::Tuple{})
+    ℓ, _, x_rest = logdensityof_with_rest(μ, x)
+    return ℓ, x_rest
 end
 
-@inline function _batched_consume(A::AbstractArray, ::Tuple{})
-    A_flat, A_rest = _batched_split(A, 1)
-    batch_axes = Base.tail(axes(A))
-    return view(A_flat, firstindex(A_flat, 1), batch_axes...), A_rest
+function _stream_ld_with_rest(f::F, μ, X::AbstractArray, sz::Dims) where {F}
+    vsz = _stream_consume_size(μ)
+    X_μ, X_rest = _batched_consume(X, vsz, sz)
+    return _consumed_ld(f, μ, X_μ, vsz), X_rest
 end
 
-@inline function _batched_split(A::AbstractArray, n::Integer)
+# Scalar variates are consumed as `(1, sz..., batch dims...)` and the leading
+# dimension is summed out, so that no rank-0 arrays arise:
+@inline _consumed_ld(f::F, μ, X_μ, ::Tuple{}) where {F} = _sum_leading_dims(_batched_kernel(f, μ, X_μ), static(1))
+@inline _consumed_ld(f::F, μ, X_μ, ::SizeLike) where {F} = _batched_kernel(f, μ, X_μ)
+
+# Consume `prod(sz)` variates of flat size `vsz` from the leading rows of a
+# batch of streams as a flat batch `(vsz..., sz..., batch dims...)`; scalar
+# variates as `(1, sz..., batch dims...)`:
+@inline function _batched_consume(X::AbstractArray, vsz::SizeLike, sz::Dims)
+    dims = _consumed_dims(vsz)
+    n_rows = prod(dims) * prod(sz)
+    X_flat, X_rest = _batched_split(X, n_rows)
+    return _reshape_consumed(X_flat, (dims..., sz...)), X_rest
+end
+@inline _consumed_dims(::Tuple{}) = (1,)
+@inline _consumed_dims(vsz::SizeLike) = map(dynamic, _size_dims(vsz))
+
+@inline _reshape_consumed(X_flat::AbstractArray, ::Tuple{Int}) = X_flat
+@inline function _reshape_consumed(X_flat::AbstractArray, dims::Tuple{Vararg{Int}})
+    reshape(X_flat, (dims..., Base.tail(size(X_flat))...))
+end
+
+@inline function _batched_split(A::AbstractArray, n_rows::Integer)
     stream_idxs = axes(A, 1)
-    if length(stream_idxs) < n
-        throw(ArgumentError("Variate streams too short during batched density evaluation"))
+    if length(stream_idxs) < n_rows
+        throw(ArgumentError("Variate streams too short during batched evaluation"))
     end
     batch_axes = Base.tail(axes(A))
     i0 = first(stream_idxs)
-    A_flat = view(A, i0:(i0 + n - 1), batch_axes...)
-    A_rest = view(A, (i0 + n):last(stream_idxs), batch_axes...)
+    A_flat = view(A, i0:(i0 + n_rows - 1), batch_axes...)
+    A_rest = view(A, (i0 + n_rows):last(stream_idxs), batch_axes...)
     return A_flat, A_rest
 end
 
-# Chunks of variates with more than one flat dimension are reshaped, the
-# batch dimensions are dynamic anyway:
-@inline _batched_chunk_shape(A_flat::AbstractArray, ::Tuple{IntegerLike}) = A_flat
-@inline function _batched_chunk_shape(A_flat::AbstractArray, dims::Tuple{Vararg{IntegerLike}})
-    reshape(A_flat, (map(dynamic, dims)..., Base.tail(size(A_flat))...))
+# Static streams split into static chunks:
+@inline function _batched_split(A::StaticVector, n_rows::StaticInteger{N}) where {N}
+    idxs = maybestatic_eachindex(A)
+    i0 = maybestatic_first(idxs)
+    _get_or_view(A, i0, i0 + n_rows - static(1)), _get_or_view(A, i0 + n_rows, maybestatic_last(idxs))
 end
 
-function _batched_consume(::AbstractArray, sz::NoMSpaceElementSize)
-    throw(ArgumentError("Batched stream consumption requires a variate of known flat size"))
+@noinline function _throw_stream_too_long()
+    throw(ArgumentError("Variate streams too long during density evaluation"))
+end
+
+# Whether all variates consumed by a measure from streams have sizes that
+# are fixed at the type level, so that batches of streams can be consumed
+# in fused operations; otherwise a batch of streams is consumed stream by
+# stream by the outermost stream combinator.
+@inline fixed_stream_size(μ::MU) where {MU} = fixed_stream_size(MU)
+@inline fixed_stream_size(::Type{MU}) where {MU} = static(mspace_ndims(MU) isa Integer)
+
+# Batches of streams consumed stream by stream (host loop):
+function _streamwise_ld(f::F, μ, X::AbstractArray) where {F}
+    map(sliced(X, Val(1))) do x
+        ℓ, x_rest = batched_logdensityof_with_rest(μ, x, ())
+        isempty(x_rest) || _throw_stream_too_long()
+        _dynamic_logd(_materialize(ℓ), x)
+    end
 end
diff --git a/src/density-core.jl b/src/density-core.jl
index c7de098b..71ef2752 100644
--- a/src/density-core.jl
+++ b/src/density-core.jl
@@ -28,7 +28,7 @@ To compute log-density relative to `basemeasure(m)` or *define* a log-density
 To compute a log-density relative to a specific base-measure, see
 `logdensity_rel`.
 """
-@inline logdensityof(μ::AbstractMeasure, x) = logdensityof_impl(μ, x)
+@inline logdensityof(μ::AbstractMeasure, x) = _point_ld(logdensityof_impl, μ, x)
 
 """
     MeasureBase.logdensityof_impl(μ::AbstractMeasure, x)
@@ -76,42 +76,32 @@ end
 """
     MeasureBase.logdensityof_with_rest(μ::AbstractMeasure, x)
 
-Compute the log-density of `μ` at the beginning of `x`, a flat stream of
-variate content that may extend beyond the variate of `μ`.
+Consume the variate of `μ` at the beginning of the flat vector stream `x`
+(or the named entries of the `NamedTuple` `x`) and compute its log-density.
 
-`x` must either be a vector that starts with the (flattened) variate of `μ`
-(for measures combined via `vcat`) or a `NamedTuple` whose first properties
-constitute the variate of `μ` (for measures combined via `merge`).
-
-Returns a tuple `(ℓ, x_μ, x_rest)` of the log-density `ℓ`, the variate
-`x_μ` of `μ` consumed from the stream, and the unconsumed rest of the
-stream.
-
-Measure types whose variate size depends on measure values, like
-[`mbind`](@ref) results, implement density calculation via
-`logdensityof_with_rest` instead of
-[`logdensityof_impl`](@ref MeasureBase.logdensityof_impl).
-
-The default implementation determines the size resp. the property names of
-the variate via [`some_mspace_elsize`](@ref MeasureBase.some_mspace_elsize)
-resp. `testvalue` and delegates to `logdensityof_impl`.
+Returns a tuple `(ℓ, x_μ, x_rest)` of the log-density, the consumed variate
+`x_μ` and the unconsumed rest of `x`. Measures whose variates have a fixed
+size consume that size (see [`MeasureBase.mspace_flatsize`](@ref) and
+[`MeasureBase.some_mspace_elsize`](@ref)), measures with variates of
+value-dependent size implement the consumption themselves. Batches of
+streams are consumed by [`MeasureBase.batched_logdensityof_with_rest`](@ref).
 """
 function logdensityof_with_rest end
 
 function logdensityof_with_rest(μ::AbstractMeasure, x::AbstractVector)
     a, x_rest = _consume_from_stream(x, _stream_consume_size(μ))
-    return logdensityof_impl(μ, a), a, x_rest
+    return _point_ld(logdensityof_impl, μ, a), a, x_rest
 end
 
-@inline _stream_consume_size(μ) = _stream_consume_size(μ, mspace_flatsize(μ))
-@inline _stream_consume_size(μ, sz::SizeLike) = sz
-@inline _stream_consume_size(μ, ::NoMSpaceElementSize) = some_mspace_elsize(μ)
-
 function logdensityof_with_rest(μ::AbstractMeasure, x::NamedTuple)
     a, x_rest = _split_after(x, Val(_mspace_names(μ)))
     return logdensityof_impl(μ, a), a, x_rest
 end
 
+@inline _stream_consume_size(μ) = _stream_consume_size(μ, mspace_flatsize(μ))
+@inline _stream_consume_size(μ, sz::SizeLike) = sz
+@inline _stream_consume_size(μ, ::NoMSpaceElementSize) = some_mspace_elsize(μ)
+
 _mspace_names(μ::AbstractMeasure) = keys(testvalue(μ))
 
 
diff --git a/src/density.jl b/src/density.jl
index 7c56f43c..361b0851 100644
--- a/src/density.jl
+++ b/src/density.jl
@@ -223,6 +223,8 @@ basemeasure(μ::DensityMeasure) = μ.base
 @inline mspace_elsize(μ::DensityMeasure) = mspace_elsize(μ.base)
 @inline mspace_flatsize(μ::DensityMeasure) = mspace_flatsize(μ.base)
 @inline mspace_flatsize(::Type{<:DensityMeasure{<:Any,B}}) where {B} = mspace_flatsize(B)
+@inline mspace_ndims(::Type{<:DensityMeasure{<:Any,B}}) where {B} = mspace_ndims(B)
+@inline fixed_stream_size(::Type{<:DensityMeasure{<:Any,B}}) where {B} = fixed_stream_size(B)
 
 logdensity_def(μ::DensityMeasure, x) = logdensityof(μ.f, x)
 
diff --git a/src/mspace.jl b/src/mspace.jl
index ca0d269e..c1520863 100644
--- a/src/mspace.jl
+++ b/src/mspace.jl
@@ -100,3 +100,30 @@ Composite measures use it to determine the flat size of their variates
 without inspecting each component.
 """
 @inline mspace_flatsize(::Type{MU}) where {MU} = NoMSpaceElementSize{MU}()
+
+
+"""
+    MeasureBase.mspace_ndims(::Type{MU})
+    MeasureBase.mspace_ndims(μ)
+
+The number of dimensions of the flat variates of measures of type `MU`,
+`0` for scalar variates, or a [`MeasureBase.NoMSpaceElementSize`](@ref)
+if unknown.
+
+Batched kernels rely on it to tell the variate dimensions of a flat batch
+from its batch dimensions. It follows from
+[`MeasureBase.mspace_flatsize`](@ref) where that is known, measure types
+with array variates of dynamic size declare it directly.
+"""
+function mspace_ndims end
+
+@inline mspace_ndims(::Type{MU}) where {MU} = _ndims_of_size(mspace_flatsize(MU), MU)
+@inline mspace_ndims(μ::MU) where {MU} = _ndims_of_size(mspace_flatsize(μ), MU, mspace_ndims(MU))
+
+@inline _ndims_of_size(sz::SizeLike, ::Type) = length(_size_dims(sz))
+@inline _ndims_of_size(::NoMSpaceElementSize, ::Type{MU}) where {MU} = NoMSpaceElementSize{MU}()
+@inline _ndims_of_size(sz::SizeLike, ::Type, ::Any) = length(_size_dims(sz))
+@inline _ndims_of_size(::NoMSpaceElementSize, ::Type, n) = n
+
+@inline _add_ndims(n::Integer, k::Integer) = n + k
+@inline _add_ndims(n::NoMSpaceElementSize, ::Integer) = n
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 1e542339..90f4bdcb 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -68,3 +68,17 @@ function batched_transport_from_std(::Type{S}, μ::Dirac, Z::AbstractArray) wher
     X .= μ.x
     return X
 end
+
+@inline mspace_ndims(::Type{<:Dirac{<:AbstractArray{<:Any,N}}}) where {N} = N
+
+# Batches of array variates: all elements of a variate must match.
+function batched_logdensityof_impl(μ::Dirac{<:AbstractArray{<:Any,N}}, X::AbstractArray) where {N}
+    matches = _all_leading_dims(X .== μ.x, static(N))
+    ifelse.(matches, zero(_logd_numtype(X)), _neg_inf_logd(X))
+end
+batched_logdensity_def(μ::Dirac{<:AbstractArray}, X::AbstractArray) = _zero_logd_batch(X, static(ndims(μ.x)))
+
+@inline _all_leading_dims(A::AbstractArray{Bool,N}, ::StaticInteger{N}) where {N} = all(A)
+@inline function _all_leading_dims(A::AbstractArray{Bool}, ::StaticInteger{N}) where {N}
+    dropdims(all(A; dims = ntuple(identity, Val(N))); dims = ntuple(identity, Val(N)))
+end
diff --git a/src/transport-batched.jl b/src/transport-batched.jl
index 5a8cedcf..3ef0e0ff 100644
--- a/src/transport-batched.jl
+++ b/src/transport-batched.jl
@@ -82,9 +82,13 @@ of the streams.
 function batched_transport_to_std_with_rest end
 
 function batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray) where {S<:StdMeasure}
-    X_μ, X_rest = _batched_consume(X, mspace_flatsize(μ))
-    return batched_transport_to_std(S, μ, X_μ), X_μ, X_rest
+    vsz = _stream_consume_size(μ)
+    X_μ, X_rest = _batched_consume(X, vsz, ())
+    X_v = _consumed_variates(X_μ, vsz)
+    return batched_transport_to_std(S, μ, X_v), X_v, X_rest
 end
+@inline _consumed_variates(X_μ::AbstractArray, ::Tuple{}) = _drop_stdstream_dim(X_μ)
+@inline _consumed_variates(X_μ::AbstractArray, ::SizeLike) = X_μ
 
 
 """
diff --git a/test/combinators/combined.jl b/test/combinators/combined.jl
index 43d9e20c..3b6e3e1c 100644
--- a/test/combinators/combined.jl
+++ b/test/combinators/combined.jl
@@ -50,7 +50,8 @@ using AffineMaps: Mul
         @test mpab isa MeasureBase.CombinedMeasure
         x = vcat(randn(2), rand(2))
         @test logdensityof(mpab, x) ≈ logdensityof(mab, x)
-        @test_throws ArgumentError logdensities(mpab, vcat(randn(2, 3), rand(2, 3)))
+        Xab = vcat(randn(2, 3), rand(2, 3))
+        @test logdensities(mpab, Xab) ≈ [logdensityof(mpab, x) for x in eachcol(Xab)]
     end
 
     @testset "CombinedMeasure" begin
diff --git a/test/logdensities.jl b/test/logdensities.jl
index 5ecf21c4..ee4f1a34 100644
--- a/test/logdensities.jl
+++ b/test/logdensities.jl
@@ -187,9 +187,11 @@ MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s)
         X = vcat(randn(2, 6), rand(3, 6))
         @test @inferred(logdensities(m, X)) ≈ [logdensityof(m, x) for x in eachcol(X)]
         @test logdensities(m, sliced(X, 1)) ≈ logdensities(m, X)
-        ℓ, A_μ, A_rest = MeasureBase.batched_logdensityof_with_rest(StdNormal()^2, X)
+        ℓ, A_rest = MeasureBase.batched_logdensityof_with_rest(StdNormal()^2, X, ())
         @test ℓ ≈ vec(sum(stdnormal_ld.(X[1:2, :]), dims = 1))
-        @test size(A_μ) == (2, 6) && size(A_rest) == (3, 6)
+        @test size(A_rest) == (3, 6)
+        ℓ2, A_rest2 = MeasureBase.batched_logdensityof_with_rest(StdNormal(), X, (2,))
+        @test ℓ2 ≈ stdnormal_ld.(X[1:2, :]) && size(A_rest2) == (3, 6)
         @test_throws ArgumentError logdensities(m, vcat(X, rand(1, 6)))
 
         m3 = mcombine(vcat, StdNormal(), mcombine(vcat, StdExponential()^2, StdLogistic()))
@@ -210,7 +212,7 @@ MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s)
         @test @inferred(logdensityof(ms, SVector{5}(x))) ≈ logdensityof(m, x)
         @test @inferred(logdensityof(ms, x)) ≈ logdensityof(m, x)
         @test logdensities(ms, X) ≈ logdensities(m, X)
-        @test_throws ArgumentError MeasureBase.batched_logdensityof_with_rest(StdNormal(), zeros(0, 4))
+        @test_throws ArgumentError MeasureBase.batched_logdensityof_with_rest(StdNormal(), zeros(0, 4), ())
     end
 
     @testset "array products of array-variate marginals" begin

From feab91db0be26e03bbc95d6feef19a7c81d9e62b Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 10:06:25 +0200
Subject: [PATCH 093/122] Store arrays of isbits marginals as struct arrays

productmeasure keeps arrays of parameterized isbits marginals as
StructArrays, unwrapping nested parameter structs into nested struct
arrays, so that the fused kernels of array products broadcast over the
numeric parameter columns: the marginals are rebuilt from the column
values inside the kernel, singleton-typed columns contribute their
instance. This makes such products work on GPU arrays and in traced
code alike. Measures holding arrays get Adapt rules, and the degrees of
freedom of scalar-variate marginal arrays are counted without reducing
over the marginals.

Created by generative AI.
---
 Project.toml                          |  4 ++
 src/MeasureBase.jl                    |  4 ++
 src/combinators/combined.jl           |  2 +
 src/combinators/power.jl              |  2 +
 src/combinators/product.jl            | 71 ++++++++++++++++++++++++---
 src/combinators/smart-constructors.jl | 21 +++++++-
 src/combinators/transformedmeasure.jl |  6 +++
 src/combinators/weighted.jl           |  2 +
 src/primitives/dirac.jl               |  2 +
 test/Project.toml                     |  2 +
 test/combinators/product.jl           | 65 ++++++++++++++++++++++++
 test/runtests.jl                      |  1 +
 12 files changed, 172 insertions(+), 10 deletions(-)
 create mode 100644 test/combinators/product.jl

diff --git a/Project.toml b/Project.toml
index e65a73b2..8398bc26 100644
--- a/Project.toml
+++ b/Project.toml
@@ -5,6 +5,7 @@ authors = ["Chad Scherrer <chad.scherrer@gmail.com>", "Oliver Schulz <oschulz@mp
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
@@ -31,6 +32,7 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 StaticThings = "7e4b4f32-fbf9-4b74-9510-4d15222ac973"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
@@ -60,6 +62,8 @@ MeasureBaseMooncakeExt = "Mooncake"
 MeasureBaseReactantExt = "Reactant"
 
 [compat]
+Adapt = "3.7, 4"
+StructArrays = "0.6.18, 0.7"
 ArgCheck = "1, 2"
 ArraysOfArrays = "1.3"
 ChainRulesCore = "1"
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index f4e87bb8..5bee3e58 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -69,6 +69,10 @@ using ArraysOfArrays:
 
 using OneTwoMany: firstarg, secondarg
 
+using StructArrays: StructArray
+import StructArrays
+import Adapt
+
 export gentype
 
 export AbstractMeasure
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index 55c92bce..5cf9f16f 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -331,3 +331,5 @@ function batched_transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure{typ
     X = vcat(_as_stream_batch(A, mspace_flatsize(μ.α)), _as_stream_batch(B, mspace_flatsize(μ.β)))
     return X, Z_rest
 end
+
+Adapt.adapt_structure(to, μ::CombinedMeasure) = mcombine(μ.f_c, Adapt.adapt(to, μ.α), Adapt.adapt(to, μ.β))
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 49f6ba00..ff74355f 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -403,3 +403,5 @@ end
     chunks = sliced(reshape(Z, (dynamic(fast_dof(ν)), n_variates, batch_dims...)), Val(1))
     stacked(map(Base.Fix1(_FromStd{S}(), ν), chunks))
 end
+
+Adapt.adapt_structure(to, μ::PowerMeasure) = PowerMeasure(Adapt.adapt(to, pwr_base(μ)), pwr_axes(μ))
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index c1f06fe7..f77dd668 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -222,8 +222,7 @@ end
 @inline function _array_product_kernel(f::F, μ::ProductMeasure, X::AbstractArray, ::StaticInteger{0}) where {F}
     mar = marginals(μ)
     _check_flatsize(X, maybestatic_size(mar))
-    ld = Broadcast.instantiate(Broadcast.broadcasted(_DynamicPointLogd(f), mar, X))
-    _sum_leading_dims(ld, static(ndims(mar)))
+    _sum_leading_dims(_marginal_broadcast(_DynamicPointLogd(f), mar, X), static(ndims(mar)))
 end
 @inline function _array_product_kernel(f::F, μ::ProductMeasure, X::AbstractArray, ::StaticInteger{K}) where {F,K}
     _marginal_slices_ld(f, marginals(μ), X, Val(K), Val(ndims(X) - K - ndims(marginals(μ))))
@@ -321,7 +320,16 @@ fast_dof(d::AbstractProductMeasure) = _sum_dofs(fast_dof, marginals(d))
 # Sums over static DOFs of tuples fold at compile time, arrays of marginals
 # are summed dynamically (also on GPU arrays):
 @inline _sum_dofs(f, mar) = sum(f, mar)
-@inline _sum_dofs(f, mar::AbstractArray) = mapreduce(_dynamic_dof ∘ f, +, mar; init = 0)
+@inline _sum_dofs(f, mar::AbstractArray{M}) where {M} = _sum_dofs(f, mar, _unit_dof(M))
+@inline _sum_dofs(f, mar::AbstractArray, ::True) = length(mar)
+@inline _sum_dofs(f, mar::AbstractArray, ::False) = mapreduce(_dynamic_dof ∘ f, +, mar; init = 0)
+
+# Marginals with scalar variates and a standard transport have one degree
+# of freedom each, so their total needs no reduction over the marginals
+# (which may live on a device):
+@inline function _unit_dof(::Type{M}) where {M}
+    static(mspace_ndims(M) === 0 && preferred_stdmeasure(M) isa Type{<:StdMeasure})
+end
 @inline _sum_dofs(f, mar::StaticArray) = mapreduce(f, +, mar; init = static(0))
 @inline _dynamic_dof(n::IntegerLike) = dynamic(n)
 @inline _dynamic_dof(nodof::AbstractNoDOF) = nodof
@@ -377,7 +385,7 @@ function transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, x::
     _array_product_to_std(S, μ, x, Val(isconcretetype(M)))
 end
 function _array_product_to_std(::Type{S}, μ, x::AbstractArray, ::Val{true}) where {S}
-    _flat_std_of(broadcast(_ToStd{S}(), marginals(μ), x))
+    _flat_std_of(_materialize(_marginal_broadcast(_ToStd{S}(), marginals(μ), x)))
 end
 # Marginals of mixed types may have standard variates of mixed shapes:
 function _array_product_to_std(::Type{S}, μ, x::AbstractArray, ::Val{false}) where {S}
@@ -410,7 +418,7 @@ function transport_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, z
 end
 function _array_product_from_std(::Type{S}, μ, z::AbstractVector, ::Tuple{}) where {S}
     mar = marginals(μ)
-    broadcast(_FromStd{S}(), mar, maybestatic_reshape(z, maybestatic_size(mar)))
+    _materialize(_marginal_broadcast(_FromStd{S}(), mar, maybestatic_reshape(z, maybestatic_size(mar))))
 end
 function _array_product_from_std(::Type{S}, μ, z::AbstractVector, ::Any) where {S}
     ys, z_rest = _marginals_from_std_with_rest(S, marginals(μ), z)
@@ -434,7 +442,7 @@ function _marginals_from_std_with_rest(::Type{S}, νs::AbstractArray{M}, z::Abst
         # The variate type is uniform, so the loop is type stable (the type
         # of the remaining stream stays invariant under repeated view-taking):
         y1, z_rest = transport_from_std_with_rest(S, νs[first(idxs)], z)
-        ys = similar(νs, typeof(y1))
+        ys = similar(Array{typeof(y1)}, axes(νs))
         ys[first(idxs)] = y1
         for i in Iterators.drop(idxs, 1)
             ys[i], z_rest = transport_from_std_with_rest(S, νs[i], z_rest)
@@ -464,7 +472,7 @@ function batched_transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{
 end
 function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Val{true}) where {S}
     _check_flatsize(X, maybestatic_size(marginals(μ)))
-    _as_stream_batch(broadcast(_ToStd{S}(), marginals(μ), X), maybestatic_size(marginals(μ)))
+    _as_stream_batch(_materialize(_marginal_broadcast(_ToStd{S}(), marginals(μ), X)), maybestatic_size(marginals(μ)))
 end
 function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Val{false}) where {S}
     _batched_to_std(S, μ, X, mspace_flatsize(μ))
@@ -475,7 +483,7 @@ function batched_transport_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArra
 end
 function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{true}) where {S}
     mar = marginals(μ)
-    broadcast(_FromStd{S}(), mar, reshape(Z, (map(dynamic, maybestatic_size(mar))..., Base.tail(size(Z))...)))
+    _materialize(_marginal_broadcast(_FromStd{S}(), mar, reshape(Z, (map(dynamic, maybestatic_size(mar))..., Base.tail(size(Z))...))))
 end
 function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{false}) where {S}
     _batched_from_std(S, μ, Z, mspace_flatsize(μ))
@@ -511,3 +519,50 @@ end
 @inline function fixed_stream_size(::Type{<:ProductMeasure{NamedTuple{names,M}}}) where {names,M<:Tuple}
     fixed_stream_size(ProductMeasure{M})
 end
+
+
+# Broadcasts over struct arrays of marginals run over their leaf columns,
+# the marginals are rebuilt from the column values inside the kernel:
+
+@inline _leaf_columns(sa::StructArray) = _leaf_columns_of(values(StructArrays.components(sa)))
+@inline _leaf_columns_of(cs::Tuple) = (_leaf_columns_of(first(cs))..., _leaf_columns_of(Base.tail(cs))...)
+@inline _leaf_columns_of(::Tuple{}) = ()
+@inline _leaf_columns_of(c::StructArray) = _leaf_columns(c)
+@inline _leaf_columns_of(c::AbstractArray{T}) where {T} = Base.issingletontype(T) ? () : (c,)
+
+@generated function _rebuild_element(::Type{SA}, vals::Tuple) where {SA<:StructArray}
+    expr, _ = _rebuild_expr(SA, 1)
+    return expr
+end
+function _rebuild_expr(::Type{SA}, i::Int) where {T,N,C,SA<:StructArray{T,N,C}}
+    args = Any[]
+    for CT in C.parameters[2].parameters
+        if CT <: StructArray
+            e, i = _rebuild_expr(CT, i)
+            push!(args, e)
+        elseif Base.issingletontype(eltype(CT))
+            push!(args, :($(eltype(CT).instance)))
+        else
+            push!(args, :(vals[$i]))
+            i += 1
+        end
+    end
+    return :(constructorof($T)($(args...))), i
+end
+
+struct _WithElement{SA,G} <: Function
+    g::G
+end
+_WithElement{SA}(g::G) where {SA,G} = _WithElement{SA,G}(g)
+@inline function (k::_WithElement{SA})(args::Vararg{Any,N}) where {SA,N}
+    k.g(_rebuild_element(SA, Base.front(args)), args[end])
+end
+
+@inline function _marginal_broadcast(g::G, mar::AbstractArray, X) where {G}
+    Broadcast.instantiate(Broadcast.broadcasted(g, mar, X))
+end
+@inline function _marginal_broadcast(g::G, mar::StructArray, X) where {G}
+    Broadcast.instantiate(Broadcast.broadcasted(_WithElement{typeof(mar)}(g), _leaf_columns(mar)..., X))
+end
+
+Adapt.adapt_structure(to, μ::ProductMeasure) = ProductMeasure(Adapt.adapt(to, marginals(μ)))
diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index be69d142..e6ba3867 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -117,12 +117,29 @@ function _generic_productmeasure_impl(mar::AbstractArray{T}) where {T}
     if Base.issingletontype(T)
         powermeasure(instance(T), axes(mar))
     elseif T <: AbstractMeasure
-        ProductMeasure(mar)
+        ProductMeasure(_marginal_storage(mar))
     else
-        ProductMeasure(map(asmeasure, mar))
+        ProductMeasure(_marginal_storage(map(asmeasure, mar)))
     end
 end
 
+# Arrays of parameterized isbits marginals are stored as struct arrays, so
+# that batched kernels broadcast over the numeric parameter columns (also
+# on GPUs and in traced code):
+function _marginal_storage(mar::AbstractArray{T}) where {T}
+    if isconcretetype(T) && isbitstype(T) && !Base.issingletontype(T)
+        StructArray(mar; unwrap = _unwrap_field)
+    else
+        mar
+    end
+end
+_marginal_storage(mar::StructArray) = mar
+
+@inline function _unwrap_field(::Type{T}) where {T}
+    isstructtype(T) && !Base.issingletontype(T) && fieldcount(T) > 0 &&
+        !(T <: Number) && !(T <: AbstractArray) && !(T <: Tuple) && !(T <: AbstractString) && !(T <: Symbol)
+end
+
 @inline function _generic_productmeasure_impl(
     mar::ReadonlyMappedArray{T,N,A,Returns{M}},
 ) where {T,N,A,M}
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 9e9682de..9d835221 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -101,6 +101,8 @@ end
     _ndims_of_size_type(VS, MU)
 end
 @inline _ndims_of_size_type(::Type{<:Tuple{Vararg{Any,N}}}, ::Type) where {N} = N
+@inline mspace_flatsize(::Type{<:PushforwardMeasure{<:Any,<:Any,<:Any,<:Any,Tuple{}}}) = ()
+@inline mspace_flatsize(::Type{<:PushforwardMeasure{<:Any,<:Any,<:Any,<:Any,StaticArrays.Size{S}}}) where {S} = StaticArrays.Size(S)
 @inline _ndims_of_size_type(::Type{StaticArrays.Size{S}}, ::Type) where {S} = length(S)
 @inline _ndims_of_size_type(::Type, ::Type{MU}) where {MU} = NoMSpaceElementSize{MU}()
 
@@ -330,3 +332,7 @@ function _pullback_impl(f, μ, style = AdaptRootMeasure())
 end
 
 @deprecate pullback(f, μ, style::PushFwdStyle = AdaptRootMeasure()) pullbck(f, μ, style)
+
+function Adapt.adapt_structure(to, ν::PushforwardMeasure)
+    PushforwardMeasure(Adapt.adapt(to, ν.f), Adapt.adapt(to, ν.finv), Adapt.adapt(to, ν.origin), ν.style, ν.varsize)
+end
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index fba7806a..d3969bcb 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -93,3 +93,5 @@ insupport(μ::WeightedMeasure, x) = insupport(μ.base, x)
     batched_transport_to_std_with_rest(S, basemeasure(μ), X)
 @inline batched_transport_from_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, Z::AbstractArray) where {S<:StdMeasure} =
     batched_transport_from_std_with_rest(S, basemeasure(μ), Z)
+
+Adapt.adapt_structure(to, μ::WeightedMeasure) = WeightedMeasure(μ.logweight, Adapt.adapt(to, μ.base))
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 90f4bdcb..9fbff0a6 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -82,3 +82,5 @@ batched_logdensity_def(μ::Dirac{<:AbstractArray}, X::AbstractArray) = _zero_log
 @inline function _all_leading_dims(A::AbstractArray{Bool}, ::StaticInteger{N}) where {N}
     dropdims(all(A; dims = ntuple(identity, Val(N))); dims = ntuple(identity, Val(N)))
 end
+
+Adapt.adapt_structure(to, μ::Dirac) = Dirac(Adapt.adapt(to, μ.x))
diff --git a/test/Project.toml b/test/Project.toml
index 81579398..7c717b0d 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,4 +1,5 @@
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 AffineMaps = "2c83c9a8-abf5-4329-a0d7-deffaf474661"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
@@ -26,6 +27,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/combinators/product.jl b/test/combinators/product.jl
new file mode 100644
index 00000000..5a4ee6e1
--- /dev/null
+++ b/test/combinators/product.jl
@@ -0,0 +1,65 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using Random, Statistics
+using StableRNGs: StableRNG
+using StructArrays: StructArray
+using Adapt: adapt
+using JLArrays
+using AffineMaps: Mul
+using ArraysOfArrays: sliced, flatview
+using InverseFunctions: inverse
+
+using MeasureBase
+using MeasureBase: StdNormal, StdUniform, productmeasure, pushfwd, marginals, transport_to, logdensities
+
+@testset "products over arrays of marginals" begin
+    stblrng() = StableRNG(789990641)
+
+    @testset "struct array storage" begin
+        P = productmeasure([pushfwd(Mul(s), StdNormal()) for s in (1.0, 2.0, 3.0)])
+        mar = marginals(P)
+        @test mar isa StructArray
+        @test length(mar) == 3 && mar[2] == pushfwd(Mul(2.0), StdNormal())
+        @test productmeasure(mar) == P
+
+        x = randn(stblrng(), 3)
+        ℓ_ref = sum(logdensityof(m, xi) for (m, xi) in zip(mar, x))
+        @test @inferred(logdensityof(P, x)) ≈ ℓ_ref
+        @test @inferred(MeasureBase.logdensity_def(P, x)) isa Real
+        X = randn(stblrng(), 3, 5)
+        @test @inferred(logdensities(P, X)) ≈ [logdensityof(P, x) for x in eachcol(X)]
+        @test logdensities(P, sliced(X, Val(1))) ≈ logdensities(P, X)
+
+        f = transport_to(StdUniform()^3, P)
+        @test inverse(f)(f(x)) ≈ x
+        @test flatview(f.(sliced(X, Val(1)))) ≈ stack(map(f, eachcol(X)))
+        @test flatview(inverse(f).(f.(sliced(X, Val(1))))) ≈ X
+
+        @test rand(stblrng(), P) isa Vector{Float64}
+        Xr = rand(stblrng(), P^100)
+        @test size(flatview(Xr)) == (3, 100)
+        @test isapprox(vec(mean(flatview(Xr), dims = 2)), zeros(3), atol = 1.0)
+        @test isapprox(vec(std(flatview(Xr), dims = 2)), [1.0, 2.0, 3.0], rtol = 0.4)
+
+        # Non-isbits marginals keep their container:
+        Pv = productmeasure([pushfwd(Mul(randn(stblrng(), 2, 2)), StdNormal()^2) for _ in 1:2])
+        @test !(marginals(Pv) isa StructArray)
+    end
+
+    @testset "device arrays" begin
+        JLArrays.allowscalar(false)
+        P = productmeasure([pushfwd(Mul(s), StdNormal()) for s in (1.0, 2.0, 3.0)])
+        Pj = adapt(JLArray, P)
+        @test marginals(Pj) isa StructArray
+        X = randn(stblrng(), 3, 5)
+        Xj = JLArray(X)
+        ℓj = logdensities(Pj, Xj)
+        @test ℓj isa JLArray && Array(ℓj) ≈ logdensities(P, X)
+        f = transport_to(StdUniform()^3, Pj)
+        Yj = f.(sliced(Xj, Val(1)))
+        @test flatview(Yj) isa JLArray
+        @test Array(flatview(Yj)) ≈ flatview(transport_to(StdUniform()^3, P).(sliced(X, Val(1))))
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 90cf12a2..a6cf4e8c 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -34,6 +34,7 @@ include("combinators/reshape.jl")
 include("combinators/implicitlymapped.jl")
 include("combinators/combined.jl")
 include("combinators/bind.jl")
+include("combinators/product.jl")
 
 include("rand.jl")
 

From 1bb302fadd40a667111fc5081ce242d46d54214b Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 11:52:17 +0200
Subject: [PATCH 094/122] Re-plumb transport batched-first

Batched transport kernels are now the primary extension point: point
transports of powers, products and combined measures are the batched
kernels at zero batch dimensions, kernel defaults route by the declared
variate rank instead of variate sizes, and the with-rest forms take a
multiplicity of variates per stream. Powers reshape by the degrees of
freedom of their base, tuple and named tuple products transport as
tuples of batches, array products of array-variate marginals run a
marginal loop over their slices of the batch, combined measures with
value-dependent component sizes go stream by stream. Pushforwards apply
their functions to flat batches, elementwise for broadcast functions
(with fused density kernels) and as column batches for affine maps via a
new AffineMaps extension. AsMeasure gets a ConstructionBase constructor
for struct array kernels.

Created by generative AI.
---
 Project.toml                          |   3 +
 ext/MeasureBaseAffineMapsExt.jl       |  45 +++++
 src/MeasureBase.jl                    |   4 +
 src/combinators/bind.jl               |  20 +++
 src/combinators/combined.jl           |  89 ++++++---
 src/combinators/power.jl              | 221 +++++++++--------------
 src/combinators/product.jl            | 216 ++++++++++++++++------
 src/combinators/transformedmeasure.jl |  80 ++++++---
 src/combinators/weighted.jl           |  10 +-
 src/density-batched.jl                |   3 +-
 src/primitives/dirac.jl               |   9 +-
 src/transport-batched.jl              | 248 +++++++++++++++++---------
 src/transport.jl                      |  30 ++--
 test/runtests.jl                      |   1 +
 test/transport_batched.jl             | 171 ++++++++++++++++++
 15 files changed, 810 insertions(+), 340 deletions(-)
 create mode 100644 ext/MeasureBaseAffineMapsExt.jl
 create mode 100644 test/transport_batched.jl

diff --git a/Project.toml b/Project.toml
index 8398bc26..182c9cdc 100644
--- a/Project.toml
+++ b/Project.toml
@@ -39,6 +39,7 @@ 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"
@@ -50,6 +51,7 @@ 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"]
@@ -63,6 +65,7 @@ MeasureBaseReactantExt = "Reactant"
 
 [compat]
 Adapt = "3.7, 4"
+AffineMaps = "0.3"
 StructArrays = "0.6.18, 0.7"
 ArgCheck = "1, 2"
 ArraysOfArrays = "1.3"
diff --git a/ext/MeasureBaseAffineMapsExt.jl b/ext/MeasureBaseAffineMapsExt.jl
new file mode 100644
index 00000000..1861b273
--- /dev/null
+++ b/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
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 5bee3e58..de4d9528 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -126,6 +126,10 @@ struct AsMeasure{T} <: AbstractMeasure
     AsMeasure{T}(obj::T) where {T} = new(obj)
 end
 
+# Struct arrays of wrapped objects rebuild elements via ConstructionBase:
+ConstructionBase.constructorof(::Type{<:AsMeasure}) = _asmeasure
+_asmeasure(obj) = AsMeasure{typeof(obj)}(obj)
+
 Base.:(==)(a::AsMeasure, b::AsMeasure) = a.obj == b.obj
 Base.isapprox(a::AsMeasure, b::AsMeasure; kwargs...) = isapprox(a.obj, b.obj; kwargs...)
 
diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl
index a2134709..e6682b67 100644
--- a/src/combinators/bind.jl
+++ b/src/combinators/bind.jl
@@ -392,3 +392,23 @@ function transport_from_std_with_rest(::Type{S}, μ::Bind, z::AbstractVector) wh
     b, z_rest = transport_from_std_with_rest(S, _get_β_a(μ, a), z2)
     return μ.f_c(a, b), z_rest
 end
+
+function transport_from_std(::Type{S}, μ::Bind, z::AbstractVector) where {S<:StdMeasure}
+    x, z_rest = transport_from_std_with_rest(S, μ, z)
+    isempty(z_rest) || _throw_std_length_mismatch()
+    return x
+end
+
+# The secondary measure depends on the primary variate, so batches of
+# streams are consumed stream by stream (by the outermost stream
+# combinator, see `fixed_stream_size`):
+function batched_transport_to_std_with_rest(::Type{S}, μ::Bind, X::AbstractArray, sz::Dims) where {S<:StdMeasure}
+    _bind_to_std_with_rest(S, μ, X, sz)
+end
+function _bind_to_std_with_rest(::Type{S}, μ::Bind, x::AbstractVector, ::Tuple{}) where {S}
+    z, _, x_rest = transport_to_std_with_rest(S, μ, x)
+    return z, x_rest
+end
+@noinline function _bind_to_std_with_rest(::Type{S}, ::Bind, ::AbstractArray, ::Dims) where {S}
+    throw(ArgumentError("Batches of variate streams containing binds must be consumed stream by stream"))
+end
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index 5cf9f16f..21185aab 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -249,11 +249,11 @@ rand_impl(ctx::GenContext, μ::CombinedMeasure) = μ.f_c(rand_impl(ctx, μ.α),
 
 # Batches of vcat-combined measures are concatenated along the streams:
 function batched_rand_impl(ctx::GenContext, μ::CombinedMeasure{typeof(vcat)}, sz::Dims)
-    _combined_batched_rand(ctx, μ, sz, mspace_flatsize(μ.α), mspace_flatsize(μ.β))
+    _combined_batched_rand(ctx, μ, sz, _static_ndims(μ.α), _static_ndims(μ.β))
 end
-function _combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::Dims, sz_a::SizeLike, sz_b::SizeLike)
-    A = _as_stream_batch(batched_rand_impl(ctx, μ.α, sz), sz_a)
-    B = _as_stream_batch(batched_rand_impl(ctx, μ.β, sz), sz_b)
+function _combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::Dims, k_a::StaticInteger, k_b::StaticInteger)
+    A = _as_stream_batch(batched_rand_impl(ctx, μ.α, sz), k_a)
+    B = _as_stream_batch(batched_rand_impl(ctx, μ.β, sz), k_b)
     return vcat(A, B)
 end
 function _combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::Dims, ::Any, ::Any)
@@ -298,38 +298,79 @@ function transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure, z::Abstrac
     return μ.f_c(a, b), z_rest
 end
 
+function transport_from_std(::Type{S}, μ::CombinedMeasure, z::AbstractVector) where {S<:StdMeasure}
+    x, z_rest = transport_from_std_with_rest(S, μ, z)
+    isempty(z_rest) || _throw_std_length_mismatch()
+    return x
+end
+
 
-# Batched transport consumes the variate parts of both component measures
-# along batches of streams:
+# Batches of vcat-combined variates are batches of streams: with fixed
+# component sizes the whole batch is consumed in fused operations,
+# otherwise stream by stream.
 
 function batched_transport_to_std(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray) where {S<:StdMeasure}
-    Z, _, X_rest = batched_transport_to_std_with_rest(S, μ, X)
-    if size(X_rest, 1) != 0
-        throw(ArgumentError("Variate streams too long during batched transport of a combined measure"))
-    end
+    _combined_batched_to_std(S, μ, X, fixed_stream_size(μ))
+end
+function _combined_batched_to_std(::Type{S}, μ::CombinedMeasure, X::AbstractArray, ::True) where {S}
+    Z, X_rest = batched_transport_to_std_with_rest(S, μ, X, ())
+    size(X_rest, 1) == 0 || _throw_stream_too_long()
     return Z
 end
+function _combined_batched_to_std(::Type{S}, μ::CombinedMeasure, x::AbstractVector, ::False) where {S}
+    _combined_batched_to_std(S, μ, x, static(true))
+end
+function _combined_batched_to_std(::Type{S}, μ::CombinedMeasure, X::AbstractArray, ::False) where {S}
+    stacked(map(x -> _combined_batched_to_std(S, μ, x, static(true)), sliced(X, Val(1))))
+end
+
+function batched_transport_to_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray, sz::Dims) where {S<:StdMeasure}
+    _combined_to_std_with_rest(S, μ, X, sz)
+end
+function _combined_to_std_with_rest(::Type{S}, μ::CombinedMeasure, X::AbstractArray, ::Tuple{}) where {S}
+    Z_a, X2 = batched_transport_to_std_with_rest(S, μ.α, X, ())
+    Z_b, X_rest = batched_transport_to_std_with_rest(S, μ.β, X2, ())
+    return vcat(Z_a, Z_b), X_rest
+end
 
-function batched_transport_to_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray) where {S<:StdMeasure}
-    Z_a, _, X2 = batched_transport_to_std_with_rest(S, μ.α, X)
-    Z_b, _, X_rest = batched_transport_to_std_with_rest(S, μ.β, X2)
-    X_μ, _ = _batched_split(X, size(X, 1) - size(X_rest, 1))
-    return vcat(Z_a, Z_b), X_μ, X_rest
+# Several variates per stream interleave the component parts, so the rows
+# of each variate are split by the fixed component sizes:
+function _combined_to_std_with_rest(::Type{S}, μ::CombinedMeasure, X::AbstractArray, sz::Dims) where {S}
+    n_rows = _fixed_stream_length(μ.α) + _fixed_stream_length(μ.β)
+    X_μ, X_rest = _batched_split(X, n_rows * prod(sz))
+    Z, _ = _combined_to_std_with_rest(S, μ, reshape(X_μ, (n_rows, sz..., Base.tail(size(X_μ))...)), ())
+    return _merge_multiplicity(Z, sz), X_rest
 end
 
 function batched_transport_from_std(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, Z::AbstractArray) where {S<:StdMeasure}
-    X, Z_rest = batched_transport_from_std_with_rest(S, μ, Z)
-    if size(Z_rest, 1) != 0
-        throw(ArgumentError("Length of standard variates doesn't match degrees of freedom of a combined measure"))
-    end
+    _combined_batched_from_std(S, μ, Z, fixed_stream_size(μ))
+end
+function _combined_batched_from_std(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, ::True) where {S}
+    X, Z_rest = batched_transport_from_std_with_rest(S, μ, Z, ())
+    size(Z_rest, 1) == 0 || _throw_std_length_mismatch()
     return X
 end
+function _combined_batched_from_std(::Type{S}, μ::CombinedMeasure, z::AbstractVector, ::False) where {S}
+    transport_from_std(S, μ, z)
+end
+function _combined_batched_from_std(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, ::False) where {S}
+    stacked(map(z -> transport_from_std(S, μ, z), sliced(Z, Val(1))))
+end
 
-function batched_transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, Z::AbstractArray) where {S<:StdMeasure}
-    A, Z2 = batched_transport_from_std_with_rest(S, μ.α, Z)
-    B, Z_rest = batched_transport_from_std_with_rest(S, μ.β, Z2)
-    X = vcat(_as_stream_batch(A, mspace_flatsize(μ.α)), _as_stream_batch(B, mspace_flatsize(μ.β)))
-    return X, Z_rest
+function batched_transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, Z::AbstractArray, sz::Dims) where {S<:StdMeasure}
+    _combined_from_std_with_rest(S, μ, Z, sz)
+end
+# Single streams yield a variate via the point protocol:
+function _combined_from_std_with_rest(::Type{S}, μ::CombinedMeasure, z::AbstractVector, ::Tuple{}) where {S}
+    transport_from_std_with_rest(S, μ, z)
+end
+function _combined_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, ::Tuple{}) where {S}
+    A, Z2 = batched_transport_from_std_with_rest(S, μ.α, Z, ())
+    B, Z_rest = batched_transport_from_std_with_rest(S, μ.β, Z2, ())
+    return vcat(_as_stream_batch(A, _static_ndims(μ.α)), _as_stream_batch(B, _static_ndims(μ.β))), Z_rest
+end
+function _combined_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, sz::Dims) where {S}
+    _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ))
 end
 
 Adapt.adapt_structure(to, μ::CombinedMeasure) = mcombine(μ.f_c, Adapt.adapt(to, μ.α), Adapt.adapt(to, μ.β))
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index ff74355f..d5e959b7 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -136,11 +136,23 @@ end
 # Batched kernels: the base kernel runs over the flat batch, the power then
 # sums the leading dimensions of the result that belong to its axes.
 @inline function _powered_kernel(f::F, μ::PowerMeasure, X) where {F}
-    _check_pwr_batch(X, mspace_flatsize(μ))
+    _check_pwr_batch(X, μ)
     _sum_leading_dims(_batched_kernel(f, pwr_base(μ), X), static(length(pwr_axes(μ))))
 end
-@inline _check_pwr_batch(X::AbstractArray, sz_flat::SizeLike) = _check_flatsize(X, sz_flat)
-@inline _check_pwr_batch(X, ::Any) = nothing
+
+# Flat batches of powers have the power dimensions after the variate
+# dimensions of the base measure (where the rank of the base is known):
+@inline _check_pwr_batch(X::AbstractArray, μ::PowerMeasure) = _check_pwr_dims(X, _static_ndims(pwr_base(μ)), _dynamic_dims(pwr_size(μ)), false)
+@inline _check_pwr_batch(::Any, ::PowerMeasure) = nothing
+@inline function _check_pwr_dims(X::AbstractArray, ::StaticInteger{K}, dims::Dims, exact::Bool) where {K}
+    n = length(dims)
+    if (exact ? ndims(X) != K + n : ndims(X) < K + n) || ntuple(i -> size(X, K + i), Val(length(dims))) != dims
+        _throw_size_mismatch()
+    end
+    return nothing
+end
+@inline _check_pwr_dims(::AbstractArray, ::NoMSpaceElementSize, ::Dims, ::Bool) = nothing
+@inline _dynamic_dims(sz::SizeLike) = map(dynamic, _size_dims(sz))
 @inline batched_logdensityof_impl(μ::PowerMeasure, X) = _powered_kernel(logdensityof_impl, μ, X)
 @inline batched_logdensity_def(μ::PowerMeasure, X) = _powered_kernel(logdensity_def, μ, X)
 
@@ -151,20 +163,16 @@ end
 @inline logdensity_def(μ::PowerMeasure, x) = _powered_point(logdensity_def, μ, x)
 
 @inline function _powered_point(f::F, μ::PowerMeasure, x::AbstractArray{<:Number}) where {F}
-    _check_pwr_variate(μ, x, mspace_flatsize(μ))
     _point_result(_materialize(_batched_kernel(f, μ, x)), μ)
 end
 @inline function _powered_point(f::F, μ::PowerMeasure, x::AbstractArray) where {F}
+    _check_pwr_shape(μ, x)
     _powered_point_nested(f, μ, x, _flat_storage(x))
 end
 @inline function _powered_point_nested(f::F, μ::PowerMeasure, x, x_flat::AbstractArray) where {F}
-    _check_pwr_variate(μ, x, mspace_flatsize(μ))
     _point_result(_materialize(_batched_kernel(f, μ, x_flat)), μ)
 end
 function _powered_point_nested(f::F, μ::PowerMeasure, x::AbstractArray, ::NoFlatStorage) where {F}
-    if maybestatic_size(x) != pwr_size(μ)
-        _throw_size_mismatch()
-    end
     ν = pwr_base(μ)
     sum(_PointLogd(f, ν), x; init = zero(_logd_numtype(x)))
 end
@@ -172,15 +180,8 @@ end
     throw(ArgumentError("Variates of powers of measures must be arrays"))
 end
 
-# Flat variates must match the flat size where it is known, nested variates
-# the power's shape:
-@inline function _check_pwr_variate(μ::PowerMeasure, x::AbstractArray{<:Number}, sz_flat::SizeLike)
-    if !_matches_flatsize(maybestatic_size(x), sz_flat) && maybestatic_size(x) != pwr_size(μ)
-        _throw_size_mismatch()
-    end
-    return nothing
-end
-@inline function _check_pwr_variate(μ::PowerMeasure, x::AbstractArray, ::Any)
+# Nested variates have the power's shape:
+@inline function _check_pwr_shape(μ::PowerMeasure, x::AbstractArray)
     if maybestatic_size(x) != pwr_size(μ)
         _throw_size_mismatch()
     end
@@ -235,17 +236,18 @@ end
 
 # Variates may be nested arrays of the power's shape or their flat storage:
 @propagate_inbounds function checked_arg(μ::PowerMeasure, x::AbstractArray{<:Any})
-    @boundscheck begin
-        sz_x = maybestatic_size(x)
-        if sz_x != pwr_size(μ) && !_matches_flatsize(sz_x, mspace_flatsize(μ))
-            _throw_size_mismatch()
-        end
-    end
+    @boundscheck _check_pwr_variate(μ, x)
     return x
 end
 
-@inline _matches_flatsize(sz_x, sz_flat::SizeLike) = Tuple(sz_x) == Tuple(sz_flat)
-@inline _matches_flatsize(sz_x, ::NoMSpaceElementSize) = false
+@inline function _check_pwr_variate(μ::PowerMeasure, x::AbstractArray)
+    if maybestatic_size(x) != pwr_size(μ)
+        _check_pwr_flat(x, _static_ndims(pwr_base(μ)), _dynamic_dims(pwr_size(μ)))
+    end
+    return nothing
+end
+@inline _check_pwr_flat(x::AbstractArray, k::StaticInteger, dims::Dims) = _check_pwr_dims(x, k, dims, true)
+@inline _check_pwr_flat(::AbstractArray, ::NoMSpaceElementSize, ::Dims) = _throw_size_mismatch()
 
 checked_arg(μ::PowerMeasure, x::Any) = _throw_size_mismatch()
 
@@ -253,81 +255,82 @@ massof(m::PowerMeasure) = massof(m.parent)^dynamic(size2length(pwr_size(m)))
 
 
 # Transport: the standard variate of a power is the flat vector of the
-# standard variates of its base measure, in the order of the flat variate
-# storage.
+# standard variates of its innermost base measure, in the order of the flat
+# variate storage. Batches transport over the flat storage `(base variate
+# dims..., power dims..., batch dims...)`.
 
-function transport_to_std(::Type{S}, μ::PowerMeasure, x::AbstractArray) where {S<:StdMeasure}
-    _pwr_to_std(S, μ, x, _flat_storage(x), mspace_flatsize(μ))
+function batched_transport_to_std(::Type{S}, μ::PowerMeasure, X::AbstractArray) where {S<:StdMeasure}
+    _check_pwr_batch(X, μ)
+    ν, n = _pwr_unwrap(μ)
+    _merge_leading_dims(batched_transport_to_std(S, ν, X), static(1) + n)
 end
 
-# Flat storage of known flat size: transport the variates of the innermost
-# base measure over the flat storage.
-function _pwr_to_std(::Type{S}, μ::PowerMeasure, x::AbstractArray, x_flat::AbstractArray, sz_flat::SizeLike) where {S}
+function batched_transport_from_std(::Type{S}, μ::PowerMeasure, Z::AbstractArray) where {S<:StdMeasure}
     ν, _ = _pwr_unwrap(μ)
-    _check_flatsize(x_flat, sz_flat)
-    _pwr_to_std_flat(S, ν, x_flat, mspace_flatsize(ν))
+    dims = _pwr_dims(μ)
+    n_rows = _batch_dims(Z)[1]
+    dof_ν = _base_dof(n_rows, prod(dims))
+    dof_ν * prod(dims) == n_rows || _throw_std_length_mismatch()
+    batched_transport_from_std(S, ν, _reshape_batch(Z, (dof_ν, dims..., Base.tail(_batch_dims(Z))...)))
 end
 
-@inline function _pwr_to_std_flat(::Type{S}, ν, x_flat::AbstractArray, ::Tuple{}) where {S}
-    _flat_std_of(broadcast(Base.Fix1(_ToStd{S}(), ν), x_flat))
-end
+# Empty powers leave the degrees of freedom of the base undetermined:
+@inline _base_dof(n_rows::IntegerLike, n_pwr::IntegerLike) = n_rows ÷ max(n_pwr, one(n_pwr))
 
-@inline function _pwr_to_std_flat(::Type{S}, ν, x_flat::AbstractArray, sz::SizeLike) where {S}
-    _flat_std_of(map(Base.Fix1(_ToStd{S}(), ν), sliced(x_flat, Val(length(sz)))))
-end
+# All power dimensions of nested powers, innermost first:
+@inline _pwr_dims(μ::PowerMeasure) = (_pwr_dims(pwr_base(μ))..., _size_dims(pwr_size(μ))...)
+@inline _pwr_dims(ν) = ()
 
-# Otherwise transport the variates of the base measure one by one:
-function _pwr_to_std(::Type{S}, μ::PowerMeasure, x::AbstractArray, ::Any, ::Any) where {S}
+# Point transport: flat variates are batches with zero batch dimensions,
+# nested variates without flat storage transport element by element.
+function transport_to_std(::Type{S}, μ::PowerMeasure, x::AbstractArray) where {S<:StdMeasure}
+    _pwr_to_std(S, μ, x, _flat_storage(x))
+end
+@inline function _pwr_to_std(::Type{S}, μ::PowerMeasure, x, x_flat::AbstractArray) where {S}
+    _single_std(batched_transport_to_std(S, μ, x_flat))
+end
+function _pwr_to_std(::Type{S}, μ::PowerMeasure, x::AbstractArray, ::NoFlatStorage) where {S}
+    _check_pwr_shape(μ, x)
     _flat_std_of(map(Base.Fix1(_ToStd{S}(), pwr_base(μ)), x))
 end
 
 function transport_from_std(::Type{S}, μ::PowerMeasure, z::AbstractVector) where {S<:StdMeasure}
-    _check_stdlength(z, fast_dof(μ))
-    _pwr_from_std(S, μ, z, mspace_flatsize(μ))
+    _pwr_variate(μ, batched_transport_from_std(S, μ, z))
 end
 
-function _pwr_from_std(::Type{S}, μ::PowerMeasure, z::AbstractVector, sz_flat::SizeLike) where {S}
-    ν, _ = _pwr_unwrap(μ)
-    _pwr_variate(μ, _pwr_from_std_flat(S, ν, z, sz_flat, mspace_flatsize(ν)))
+# Streams: a power consumes the variates of its base measure with its size
+# as additional multiplicity. Bases without fixed variate sizes are
+# consumed element by element, for single streams.
+function batched_transport_to_std_with_rest(::Type{S}, μ::PowerMeasure, X::AbstractArray, sz::Dims) where {S<:StdMeasure}
+    _pwr_to_std_with_rest(S, μ, X, sz, fixed_stream_size(pwr_base(μ)))
 end
-
-# Base measures of unknown variate size are transported one by one:
-function _pwr_from_std(::Type{S}, μ::PowerMeasure, z::AbstractVector, ::NoMSpaceElementSize) where {S}
-    ys, z_rest = _marginals_from_std_with_rest(S, marginals(μ), z)
-    if !isempty(z_rest)
-        throw(ArgumentError("Length of standard variate doesn't match degrees of freedom of power measure"))
-    end
-    return ys
+@inline function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, X::AbstractArray, sz::Dims, ::True) where {S}
+    batched_transport_to_std_with_rest(S, pwr_base(μ), X, (_dynamic_dims(pwr_size(μ))..., sz...))
 end
-
-@inline function _check_stdlength(z::AbstractVector, n::IntegerLike)
-    if maybestatic_length(z) != n
-        throw(ArgumentError("Length of standard variate doesn't match degrees of freedom of measure"))
-    end
-    return nothing
+function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::Tuple{}, ::False) where {S}
+    z, _, x_rest = transport_to_std_with_rest(S, μ, x)
+    return z, x_rest
 end
-@inline _check_stdlength(::AbstractVector, ::AbstractNoDOF) = nothing
-
-@inline function _pwr_from_std_flat(::Type{S}, ν, z::AbstractVector, sz_flat, ::Tuple{}) where {S}
-    maybestatic_reshape(broadcast(Base.Fix1(_FromStd{S}(), ν), z), sz_flat)
+@noinline function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, ::AbstractArray, ::Dims, ::False) where {S}
+    throw(ArgumentError("Batches of variate streams containing powers of measures of type $(nameof(typeof(pwr_base(μ)))) must be consumed stream by stream"))
 end
 
-@inline function _pwr_from_std_flat(::Type{S}, ν, z::AbstractVector, sz_flat, sz_ν::SizeLike) where {S}
-    n_variates = size2length(sz_flat) ÷ size2length(sz_ν)
-    maybestatic_reshape(stacked(_pwr_from_std_chunks(S, ν, z, n_variates, fast_dof(ν))), sz_flat)
+function transport_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector) where {S<:StdMeasure}
+    _pwr_point_to_std_with_rest(S, μ, x, fixed_stream_size(pwr_base(μ)))
 end
-
-function _pwr_from_std_chunks(::Type{S}, ν, z::AbstractVector, n_variates, dof_ν::IntegerLike) where {S}
-    chunks = sliced(maybestatic_reshape(z, (dof_ν, n_variates)), Val(1))
-    map(Base.Fix1(_FromStd{S}(), ν), chunks)
+function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::True) where {S}
+    x_μ, x_rest = _consume_from_stream(x, _stream_consume_size(μ))
+    return _as_stdstream(transport_to_std(S, μ, x_μ)), x_μ, x_rest
 end
-
-function _pwr_from_std_chunks(::Type{S}, ν, z::AbstractVector, n_variates, ::AbstractNoDOF) where {S}
-    ys, z_rest = _marginals_from_std_with_rest(S, FillArrays.Fill(ν, n_variates), z)
-    if !isempty(z_rest)
-        throw(ArgumentError("Length of standard variate doesn't match degrees of freedom of power measure"))
+function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::False) where {S}
+    ν = pwr_base(μ)
+    zs = Vector{Any}(undef, length(marginals(μ)))
+    x_rest = x
+    for i in eachindex(zs)
+        zs[i], _, x_rest = transport_to_std_with_rest(S, ν, x_rest)
     end
-    return ys
+    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    return reduce(vcat, [z for z in zs]), x_μ, x_rest
 end
 
 # Powers of measures without fast degrees of freedom transport their
@@ -342,66 +345,12 @@ end
 
 # The nested variate layout of a power over its flat storage:
 @inline _pwr_variate(μ::PowerMeasure, A::AbstractArray) = _pwr_nest(pwr_base(μ), _pwr_variate(pwr_base(μ), A))
-@inline _pwr_variate(ν, A::AbstractArray) = _nest_leaf(A, mspace_flatsize(ν))
-@inline _nest_leaf(A::AbstractArray, ::Tuple{}) = A
+@inline _pwr_variate(ν, A::AbstractArray) = _nest_leaf(A, _static_ndims(ν))
+@inline _nest_leaf(A::AbstractArray, ::StaticInteger{0}) = A
 @inline _nest_leaf(A::AbstractArray, ::NoMSpaceElementSize) = A
-@inline _nest_leaf(A::AbstractArray{<:Any,N}, sz::SizeLike) where {N} = _nest_leaf(A, Val(length(sz)), Val(N))
-@inline _nest_leaf(A::AbstractArray, ::Val{N}, ::Val{N}) where {N} = A
-@inline _nest_leaf(A::AbstractArray, ::Val{M}, ::Val) where {M} = sliced(A, Val(M))
+@inline _nest_leaf(A::AbstractArray{<:Any,N}, ::StaticInteger{N}) where {N} = A
+@inline _nest_leaf(A::AbstractArray, ::StaticInteger{K}) where {K} = sliced(A, Val(K))
 @inline _pwr_nest(ν::PowerMeasure, B::AbstractArray) = sliced(B, Val(length(pwr_axes(ν))))
 @inline _pwr_nest(ν, B::AbstractArray) = B
 
-# Batched transport over the flat storage `(base variate dims..., power
-# dims..., batch dims...)`:
-
-function batched_transport_to_std(::Type{S}, μ::PowerMeasure, X::AbstractArray) where {S<:StdMeasure}
-    _pwr_batched_to_std(S, μ, X, mspace_flatsize(μ))
-end
-
-function _pwr_batched_to_std(::Type{S}, μ::PowerMeasure, X::AbstractArray, sz_flat::SizeLike) where {S}
-    ν, _ = _pwr_unwrap(μ)
-    _check_flatsize(X, sz_flat)
-    n_flat = length(sz_flat)
-    Z = _pwr_batched_to_std_flat(S, ν, X, mspace_flatsize(ν))
-    batch_dims = ntuple(i -> size(X, n_flat + i), Val(ndims(X) - n_flat))
-    return reshape(Z, (dynamic(fast_dof(μ)), batch_dims...))
-end
-
-function _pwr_batched_to_std(::Type{S}, μ::PowerMeasure, ::AbstractArray, ::NoMSpaceElementSize) where {S}
-    throw(ArgumentError("Batched transport of powers of measures of type $(nameof(typeof(pwr_base(μ)))) requires a known variate size"))
-end
-
-@inline function _pwr_batched_to_std_flat(::Type{S}, ν, X::AbstractArray, ::Tuple{}) where {S}
-    broadcast(Base.Fix1(_ToStd{S}(), ν), X)
-end
-
-@inline function _pwr_batched_to_std_flat(::Type{S}, ν, X::AbstractArray, sz::SizeLike) where {S}
-    stacked(map(Base.Fix1(_ToStd{S}(), ν), sliced(X, Val(length(sz)))))
-end
-
-function batched_transport_from_std(::Type{S}, μ::PowerMeasure, Z::AbstractArray) where {S<:StdMeasure}
-    _pwr_batched_from_std(S, μ, Z, mspace_flatsize(μ))
-end
-
-function _pwr_batched_from_std(::Type{S}, μ::PowerMeasure, Z::AbstractArray, sz_flat::SizeLike) where {S}
-    ν, _ = _pwr_unwrap(μ)
-    batch_dims = Base.tail(size(Z))
-    X = _pwr_batched_from_std_flat(S, ν, Z, batch_dims, mspace_flatsize(ν))
-    return reshape(X, (map(dynamic, _size_dims(sz_flat))..., batch_dims...))
-end
-
-function _pwr_batched_from_std(::Type{S}, μ::PowerMeasure, ::AbstractArray, ::NoMSpaceElementSize) where {S}
-    throw(ArgumentError("Batched transport to powers of measures of type $(nameof(typeof(pwr_base(μ)))) requires a known variate size"))
-end
-
-@inline function _pwr_batched_from_std_flat(::Type{S}, ν, Z::AbstractArray, batch_dims, ::Tuple{}) where {S}
-    broadcast(Base.Fix1(_FromStd{S}(), ν), Z)
-end
-
-@inline function _pwr_batched_from_std_flat(::Type{S}, ν, Z::AbstractArray, batch_dims, sz::SizeLike) where {S}
-    n_variates = size(Z, 1) ÷ dynamic(fast_dof(ν))
-    chunks = sliced(reshape(Z, (dynamic(fast_dof(ν)), n_variates, batch_dims...)), Val(1))
-    stacked(map(Base.Fix1(_FromStd{S}(), ν), chunks))
-end
-
 Adapt.adapt_structure(to, μ::PowerMeasure) = PowerMeasure(Adapt.adapt(to, pwr_base(μ)), pwr_axes(μ))
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index f77dd668..2adcb3b7 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -339,27 +339,27 @@ function checked_arg(μ::ProductMeasure{<:NTuple{N,Any}}, x::NTuple{N,Any}) wher
 end
 
 # Variates of array products are arrays of marginal variates or, for
-# marginals with array variates of known size, their flat storage:
+# marginals with array variates of declared rank, their flat storage:
 @propagate_inbounds function checked_arg(μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray) where {M}
-    @boundscheck _check_product_arg(marginals(μ), x, mspace_flatsize(M))
+    @boundscheck _check_product_arg(marginals(μ), x, _static_ndims_of(mspace_ndims(M)))
     return x
 end
 
-@inline _check_product_arg(mar, x::AbstractArray, ::Tuple{}) = _check_marginal_count(mar, x)
-@inline _check_product_arg(mar, x::AbstractArray{<:Number}, ::NoMSpaceElementSize) = _check_marginal_count(mar, x)
-@inline function _check_product_arg(mar, x::AbstractArray, ::NoMSpaceElementSize)
+@inline function _check_product_arg(mar, x::AbstractArray, ::Any)
     _check_marginal_count(mar, x)
     foreach(checked_arg, mar, x)
     return nothing
 end
-@inline function _check_product_arg(mar, x::AbstractArray, sz_m::SizeLike)
-    if size(x) == size(mar)
-        foreach(checked_arg, mar, x)
-    elseif size(x) != (Tuple(sz_m)..., size(mar)...)
+@inline function _check_product_arg(mar::AbstractArray{<:Any,N}, x::AbstractArray{<:Number}, ::StaticInteger{0}) where {N}
+    _check_marginal_count(mar, x)
+end
+@inline function _check_product_arg(mar::AbstractArray{<:Any,N}, x::AbstractArray{<:Number}, ::StaticInteger{K}) where {N,K}
+    if ndims(x) != K + N || ntuple(i -> size(x, K + i), Val(N)) != size(mar)
         _throw_marginal_mismatch()
     end
     return nothing
 end
+@inline _check_product_arg(mar, x::AbstractArray{<:Number}, ::NoMSpaceElementSize) = _check_marginal_count(mar, x)
 
 
 function checked_arg(
@@ -371,7 +371,8 @@ end
 
 
 # Transport marginal by marginal, the standard variates of the marginals
-# are concatenated in order:
+# are concatenated in order. Batches of tuple and named tuple variates are
+# tuples resp. named tuples of batches.
 
 function transport_to_std(::Type{S}, μ::ProductMeasure{<:Tuple}, x::Tuple) where {S<:StdMeasure}
     _flatten_to_rv(map((m, xi) -> _as_stdstream(transport_to_std(S, m, xi)), marginals(μ), x))
@@ -381,16 +382,21 @@ function transport_to_std(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, x:
     transport_to_std(S, productmeasure(values(marginals(μ))), values(x))
 end
 
-function transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray) where {S<:StdMeasure,M}
-    _array_product_to_std(S, μ, x, Val(isconcretetype(M)))
+function batched_transport_to_std(::Type{S}, μ::ProductMeasure{<:Tuple}, X::Tuple) where {S<:StdMeasure}
+    _vcat_std(map((m, Xi) -> batched_transport_to_std(S, m, Xi), marginals(μ), X))
 end
-function _array_product_to_std(::Type{S}, μ, x::AbstractArray, ::Val{true}) where {S}
-    _flat_std_of(_materialize(_marginal_broadcast(_ToStd{S}(), marginals(μ), x)))
+
+function batched_transport_to_std(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, X::NamedTuple{names}) where {S<:StdMeasure,names}
+    batched_transport_to_std(S, productmeasure(values(marginals(μ))), values(X))
 end
-# Marginals of mixed types may have standard variates of mixed shapes:
-function _array_product_to_std(::Type{S}, μ, x::AbstractArray, ::Val{false}) where {S}
-    zs = [_as_stdstream(transport_to_std(S, m, xi)) for (m, xi) in zip(marginals(μ), x)]
-    isempty(zs) ? SVector{0,Bool}() : reduce(vcat, zs)
+
+@inline _vcat_std(Zs::Tuple) = vcat(Zs...)
+@inline _vcat_std(::Tuple{}) = SVector{0,Bool}()
+
+function transport_from_std(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, z::AbstractVector) where {S<:StdMeasure}
+    x, z_rest = transport_from_std_with_rest(S, μ, z)
+    isempty(z_rest) || _throw_std_length_mismatch()
+    return x
 end
 
 function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:Tuple}, z::AbstractVector) where {S<:StdMeasure}
@@ -402,32 +408,38 @@ function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:NamedTuple
     return NamedTuple{names}(ys), z_rest
 end
 
-function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray}, z::AbstractVector) where {S<:StdMeasure}
-    _array_product_from_std_with_rest(S, μ, z, fast_dof(μ))
+function batched_transport_from_std(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, Z::AbstractArray) where {S<:StdMeasure}
+    X, Z_rest = batched_transport_from_std_with_rest(S, μ, Z, ())
+    size(Z_rest, 1) == 0 || _throw_std_length_mismatch()
+    return X
 end
-@inline function _array_product_from_std_with_rest(::Type{S}, μ, z, n::IntegerLike) where {S}
-    _from_std_with_rest_bydof(S, μ, z, n)
+
+function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:Tuple}, Z::AbstractArray, sz::Dims) where {S<:StdMeasure}
+    _tuple_product_from_std_with_rest(S, μ, Z, sz)
 end
-function _array_product_from_std_with_rest(::Type{S}, μ, z, ::AbstractNoDOF) where {S}
-    _marginals_from_std_with_rest(S, marginals(μ), z)
+
+function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, Z::AbstractArray, sz::Dims) where {S<:StdMeasure,names}
+    Xs, Z_rest = _tuple_product_from_std_with_rest(S, productmeasure(values(marginals(μ))), Z, sz)
+    return NamedTuple{names}(Xs), Z_rest
 end
 
-# Marginals with scalar variates transport in a single broadcast:
-function transport_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, z::AbstractVector) where {S<:StdMeasure,M}
-    _array_product_from_std(S, μ, z, mspace_flatsize(M))
+# One variate per stream is consumed marginal by marginal, several per
+# stream via the degrees of freedom of the whole product:
+function _tuple_product_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, ::Tuple{}) where {S}
+    _marginals_batched_from_std_with_rest(S, marginals(μ), Z)
 end
-function _array_product_from_std(::Type{S}, μ, z::AbstractVector, ::Tuple{}) where {S}
-    mar = marginals(μ)
-    _materialize(_marginal_broadcast(_FromStd{S}(), mar, maybestatic_reshape(z, maybestatic_size(mar))))
+function _tuple_product_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::Dims) where {S}
+    _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ))
 end
-function _array_product_from_std(::Type{S}, μ, z::AbstractVector, ::Any) where {S}
-    ys, z_rest = _marginals_from_std_with_rest(S, marginals(μ), z)
-    if !isempty(z_rest)
-        throw(ArgumentError("Length of standard variate doesn't match degrees of freedom of product measure"))
-    end
-    return ys
+
+function _marginals_batched_from_std_with_rest(::Type{S}, νs::Tuple{Vararg{Any}}, Z::AbstractArray) where {S}
+    X1, Z_rest = batched_transport_from_std_with_rest(S, νs[1], Z, ())
+    X2_end, Z_final_rest = _marginals_batched_from_std_with_rest(S, Base.tail(νs), Z_rest)
+    return (X1, X2_end...), Z_final_rest
 end
 
+_marginals_batched_from_std_with_rest(::Type{S}, ::Tuple{}, Z::AbstractArray) where {S} = (), Z
+
 function _marginals_from_std_with_rest(::Type{S}, νs::Tuple{Vararg{Any}}, z::AbstractVector) where {S}
     y1, z_rest = transport_from_std_with_rest(S, νs[1], z)
     y2_end, z_final_rest = _marginals_from_std_with_rest(S, Base.tail(νs), z_rest)
@@ -458,35 +470,131 @@ function _marginals_from_std_with_rest(::Type{S}, νs::AbstractArray{M}, z::Abst
     end
 end
 
-# Batched transport of array products with scalar-variate marginals in one
-# broadcast, the marginals align with the leading dimension of the batch:
 
-# Marginals of concrete type with scalar variates and a standard transport
-# have one degree of freedom each, so the batch aligns with the marginals:
-@inline function _fused_marginals(::Type{M}) where {M}
-    Val(isconcretetype(M) && mspace_flatsize(M) === () && preferred_stdmeasure(M) isa Type{<:StdMeasure})
-end
+# Array products: marginals of scalar variates with one degree of freedom
+# each transport in a single broadcast (the marginals align with the leading
+# dimensions of the batch), other marginals one by one over their slices of
+# the batch.
+
+@inline _fused_marginals(::Type{M}) where {M} = Val(isconcretetype(M) && _unit_dof(M) === static(true))
 
 function batched_transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, X::AbstractArray) where {S<:StdMeasure,M}
-    _array_product_batched_to_std(S, μ, X, _fused_marginals(M))
+    _array_product_batched_to_std(S, μ, X, _fused_marginals(M), _static_ndims_of(mspace_ndims(M)))
+end
+function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Val{true}, ::Any) where {S}
+    mar = marginals(μ)
+    _check_flatsize(X, maybestatic_size(mar))
+    _as_stream_batch(_materialize(_marginal_broadcast(_ToStd{S}(), mar, X)), static(ndims(mar)))
+end
+function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Val{false}, ::StaticInteger{K}) where {S,K}
+    mar = marginals(μ)
+    n_batch = ndims(X) - K - ndims(mar)
+    n_batch >= 0 || _throw_size_mismatch()
+    _marginals_to_std_loop(S, mar, X, Val(K), Val(n_batch))
 end
-function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Val{true}) where {S}
-    _check_flatsize(X, maybestatic_size(marginals(μ)))
-    _as_stream_batch(_materialize(_marginal_broadcast(_ToStd{S}(), marginals(μ), X)), maybestatic_size(marginals(μ)))
+@noinline function _array_product_batched_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, ::AbstractArray, ::Val{false}, ::NoMSpaceElementSize) where {S,M}
+    throw(ArgumentError("Batched transport of products over arrays of marginals of type $(nameof(M)) requires MeasureBase.mspace_ndims to be declared for that type"))
 end
-function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Val{false}) where {S}
-    _batched_to_std(S, μ, X, mspace_flatsize(μ))
+
+function _marginals_to_std_loop(::Type{S}, mar::AbstractArray{<:Any,N}, X::AbstractArray, ::Val{K}, ::Val{B}) where {S,N,K,B}
+    ntuple(i -> size(X, K + i), Val(N)) == size(mar) || _throw_size_mismatch()
+    lead = ntuple(_ -> Colon(), Val(K))
+    trail = ntuple(_ -> Colon(), Val(B))
+    zs = map(i -> batched_transport_to_std(S, mar[i], view(X, lead..., Tuple(i)..., trail...)), vec(CartesianIndices(mar)))
+    isempty(zs) ? similar(X, (0, ntuple(i -> size(X, K + N + i), Val(B))...)) : reduce(vcat, zs)
 end
 
 function batched_transport_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, Z::AbstractArray) where {S<:StdMeasure,M}
-    _array_product_batched_from_std(S, μ, Z, _fused_marginals(M))
+    _array_product_batched_from_std(S, μ, Z, _fused_marginals(M), _static_ndims_of(mspace_ndims(M)))
 end
-function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{true}) where {S}
+function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{true}, ::Any) where {S}
     mar = marginals(μ)
-    _materialize(_marginal_broadcast(_FromStd{S}(), mar, reshape(Z, (map(dynamic, maybestatic_size(mar))..., Base.tail(size(Z))...))))
+    size(Z, 1) == length(mar) || _throw_std_length_mismatch()
+    _materialize(_marginal_broadcast(_FromStd{S}(), mar, _reshape_batch(Z, (_batch_dims(mar)..., Base.tail(_batch_dims(Z))...))))
+end
+function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{false}, ::StaticInteger{K}) where {S,K}
+    X, Z_rest = _marginals_from_std_loop(S, marginals(μ), Z, Val(K))
+    size(Z_rest, 1) == 0 || _throw_std_length_mismatch()
+    return X
 end
-function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{false}) where {S}
-    _batched_from_std(S, μ, Z, mspace_flatsize(μ))
+@noinline function _array_product_batched_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, ::AbstractArray, ::Val{false}, ::NoMSpaceElementSize) where {S,M}
+    throw(ArgumentError("Batched transport to products over arrays of marginals of type $(nameof(M)) requires MeasureBase.mspace_ndims to be declared for that type"))
+end
+
+function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, Z::AbstractArray, sz::Dims) where {S<:StdMeasure,M}
+    _array_product_batched_from_std_with_rest(S, μ, Z, sz, _fused_marginals(M), _static_ndims_of(mspace_ndims(M)))
+end
+function _array_product_batched_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::Dims, ::Val{true}, ::Any) where {S}
+    _batched_from_std_bydof(S, μ, Z, sz, length(marginals(μ)))
+end
+function _array_product_batched_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, ::Tuple{}, ::Val{false}, ::StaticInteger{K}) where {S,K}
+    _marginals_from_std_loop(S, marginals(μ), Z, Val(K))
+end
+function _array_product_batched_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::Dims, ::Val{false}, ::StaticInteger{K}) where {S,K}
+    _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ))
+end
+@noinline function _array_product_batched_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, ::AbstractArray, ::Dims, ::Val{false}, ::NoMSpaceElementSize) where {S,M}
+    throw(ArgumentError("Batched transport to products over arrays of marginals of type $(nameof(M)) requires MeasureBase.mspace_ndims to be declared for that type"))
+end
+
+# The marginals consume the streams one after the other, their variates
+# fill the batch `(marginal variate dims..., product dims..., batch dims...)`:
+function _marginals_from_std_loop(::Type{S}, mar::AbstractArray{<:Any,N}, Z::AbstractArray, ::Val{K}) where {S,N,K}
+    idxs = vec(CartesianIndices(mar))
+    batch_dims = Base.tail(size(Z))
+    lead = ntuple(_ -> Colon(), Val(K))
+    trail = ntuple(_ -> Colon(), Val(length(batch_dims)))
+    if isempty(idxs)
+        return similar(Z, (ntuple(_ -> 0, Val(K))..., size(mar)..., batch_dims...)), Z
+    end
+    X1, Z_rest = batched_transport_from_std_with_rest(S, mar[idxs[1]], Z, ())
+    X = similar(Z, eltype(X1), (ntuple(i -> size(X1, i), Val(K))..., size(mar)..., batch_dims...))
+    X[lead..., Tuple(idxs[1])..., trail...] = X1
+    for i in idxs[2:end]
+        Xi, Z_rest = batched_transport_from_std_with_rest(S, mar[i], Z_rest, ())
+        X[lead..., Tuple(i)..., trail...] = Xi
+    end
+    return X, Z_rest
+end
+
+# Point transport of array products: arrays of marginal variates with flat
+# storage and flat variates go through the batched kernels, marginals
+# without a declared variate rank transport one by one.
+function transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray) where {S<:StdMeasure,M}
+    _array_product_to_std(S, μ, x, _flat_storage(x), _static_ndims_of(mspace_ndims(M)))
+end
+@inline function _array_product_to_std(::Type{S}, μ, x::AbstractArray, x_flat::AbstractArray, ::StaticInteger) where {S}
+    _single_std(batched_transport_to_std(S, μ, x_flat))
+end
+function _array_product_to_std(::Type{S}, μ, x::AbstractArray, ::Any, ::Any) where {S}
+    _check_marginal_count(marginals(μ), x)
+    zs = [_as_stdstream(transport_to_std(S, m, xi)) for (m, xi) in zip(marginals(μ), x)]
+    isempty(zs) ? SVector{0,Bool}() : reduce(vcat, zs)
+end
+
+# Marginals with variates of fixed size and declared rank yield a nested
+# view of the flat variate batch, others transport marginal by marginal:
+function transport_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, z::AbstractVector) where {S<:StdMeasure,M}
+    _array_product_from_std(S, μ, z, fixed_stream_size(M), _static_ndims_of(mspace_ndims(M)))
+end
+@inline function _array_product_from_std(::Type{S}, μ, z::AbstractVector, ::True, k::StaticInteger) where {S}
+    _nest_leaf(batched_transport_from_std(S, μ, z), k)
+end
+function _array_product_from_std(::Type{S}, μ, z::AbstractVector, ::Any, ::Any) where {S}
+    ys, z_rest = _marginals_from_std_with_rest(S, marginals(μ), z)
+    isempty(z_rest) || _throw_std_length_mismatch()
+    return ys
+end
+
+function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, z::AbstractVector) where {S<:StdMeasure,M}
+    _array_product_from_std_with_rest(S, μ, z, fixed_stream_size(M), _static_ndims_of(mspace_ndims(M)))
+end
+function _array_product_from_std_with_rest(::Type{S}, μ, z::AbstractVector, ::True, k::StaticInteger) where {S}
+    X, z_rest = batched_transport_from_std_with_rest(S, μ, z, ())
+    return _nest_leaf(X, k), z_rest
+end
+function _array_product_from_std_with_rest(::Type{S}, μ, z::AbstractVector, ::Any, ::Any) where {S}
+    _marginals_from_std_with_rest(S, marginals(μ), z)
 end
 
 
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 9d835221..72719767 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -106,6 +106,11 @@ end
 @inline _ndims_of_size_type(::Type{StaticArrays.Size{S}}, ::Type) where {S} = length(S)
 @inline _ndims_of_size_type(::Type, ::Type{MU}) where {MU} = NoMSpaceElementSize{MU}()
 
+# Pushforwards by elementwise functions keep the variate rank of their
+# origin:
+const _ElementwisePushfwd{M,S} = PushforwardMeasure{<:Base.BroadcastFunction,<:Base.BroadcastFunction,M,S}
+@inline mspace_ndims(::Type{MU}) where {M,MU<:_ElementwisePushfwd{M}} = mspace_ndims(M)
+
 const _NonBijectivePusfwdMeasure{M<:PushforwardMeasure,S<:PushFwdStyle} = Union{
     PushforwardMeasure{<:Any,<:NoInverse,M,S},
     PushforwardMeasure{<:NoInverse,<:Any,M,S},
@@ -184,6 +189,36 @@ for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :log
     end
 end
 
+# Pushforwards by elementwise functions evaluate densities over flat
+# batches, the log-abs-det-Jacobian terms sum over the variate dimensions
+# of the origin:
+for (bhead, head) in [(:batched_logdensityof_impl, :logdensityof_impl), (:batched_logdensity_def, :logdensity_def)]
+    @eval function $bhead(ν::_ElementwisePushfwd{M,<:AdaptRootMeasure}, Y) where {M}
+        _elementwise_pushfwd_ld($head, ν, Y, _static_ndims(ν.origin))
+    end
+    @eval function $bhead(ν::_ElementwisePushfwd{M,<:PushfwdRootMeasure}, Y) where {M}
+        _batched_kernel($head, ν.origin, broadcast(ν.finv.f, Y))
+    end
+end
+
+function _elementwise_pushfwd_ld(f::F, ν::PushforwardMeasure, Y, k::StaticInteger) where {F}
+    f_inv = ν.finv.f
+    ℓ = _batched_kernel(f, ν.origin, broadcast(f_inv, Y))
+    ladj = _sum_leading_dims(broadcast(_LadjOf(f_inv), Y), k)
+    return _lazy_combine_ladj(ℓ, ladj)
+end
+function _elementwise_pushfwd_ld(f::F, ν::PushforwardMeasure, Y, ::NoMSpaceElementSize) where {F}
+    _default_batched_kernel(f, ν, Y, _static_ndims(ν))
+end
+
+struct _LadjOf{F} <: Function
+    f::F
+end
+@inline (k::_LadjOf)(y) = last(with_logabsdet_jacobian(k.f, y))
+
+@inline _lazy_combine_ladj(ℓ::Number, ladj::Number) = _combine_logd_with_ladj(ℓ, ladj)
+@inline _lazy_combine_ladj(ℓ, ladj) = Broadcast.instantiate(Broadcast.broadcasted(_combine_logd_with_ladj, ℓ, ladj))
+
 # Checking insupport via the origin would require a potentially costly
 # transformation of x:
 insupport(m::PushforwardMeasure, x) = NoFastInsupport{typeof(m)}()
@@ -225,26 +260,27 @@ _pushfwd_dof(::Type{MU}, ::Type{<:Tuple{Any,Real}}, dof) where {MU} = dof
     return ν.f(x), z_rest
 end
 
-# Batched transport for pushforwards of measures with scalar variates, the
-# functions apply elementwise then:
-function batched_transport_to_std(::Type{S}, ν::PushforwardMeasure, Y::AbstractArray) where {S<:StdMeasure}
-    _pushfwd_batched_to_std(S, ν, Y, mspace_flatsize(ν.origin))
-end
-@inline function _pushfwd_batched_to_std(::Type{S}, ν::PushforwardMeasure, Y::AbstractArray, ::Tuple{}) where {S}
-    batched_transport_to_std(S, ν.origin, broadcast(ν.finv, Y))
+# Batches of pushforwards apply the functions to flat batches of the
+# origin, elementwise for `Base.BroadcastFunction`s and variate by variate
+# (in a host loop) otherwise. The AffineMaps extension adds affine maps.
+function batched_transport_to_std(::Type{S}, ν::PushforwardMeasure, Y) where {S<:StdMeasure}
+    batched_transport_to_std(S, ν.origin, _apply_batched(ν.finv, Y, _static_ndims(ν)))
 end
-@inline function _pushfwd_batched_to_std(::Type{S}, ν::PushforwardMeasure, Y::AbstractArray, ::Any) where {S}
-    _batched_to_std(S, ν, Y, mspace_flatsize(ν))
-end
-
 function batched_transport_from_std(::Type{S}, ν::PushforwardMeasure, Z::AbstractArray) where {S<:StdMeasure}
-    _pushfwd_batched_from_std(S, ν, Z, mspace_flatsize(ν.origin))
+    _apply_batched(ν.f, batched_transport_from_std(S, ν.origin, Z), _static_ndims(ν.origin))
 end
-@inline function _pushfwd_batched_from_std(::Type{S}, ν::PushforwardMeasure, Z::AbstractArray, ::Tuple{}) where {S}
-    broadcast(ν.f, batched_transport_from_std(S, ν.origin, Z))
-end
-@inline function _pushfwd_batched_from_std(::Type{S}, ν::PushforwardMeasure, Z::AbstractArray, ::Any) where {S}
-    _batched_from_std(S, ν, Z, mspace_flatsize(ν))
+
+# Apply `f` to a flat batch of variates of rank `k`:
+@inline _apply_batched(f, X, k) = _apply_generic(unwrap(f), X, k)
+@inline _apply_generic(f, X, k) = _apply_by_rank(f, X, k)
+@inline _apply_generic(f::Base.BroadcastFunction, X, k) = broadcast(f.f, X)
+@inline _apply_by_rank(f, X, ::StaticInteger{0}) = broadcast(f, X)
+@inline _apply_by_rank(f, X::AbstractArray, ::StaticInteger{0}) = broadcast(f, X)
+@inline _apply_by_rank(f, X::AbstractArray, ::StaticInteger{K}) where {K} = _apply_to_slices(f, X, Val(K))
+@inline _apply_to_slices(f, X::AbstractArray{<:Any,K}, ::Val{K}) where {K} = f(X)
+@inline _apply_to_slices(f, X::AbstractArray, ::Val{K}) where {K} = stacked(map(f, sliced(X, Val(K))))
+@noinline function _apply_by_rank(f, X, ::NoMSpaceElementSize)
+    throw(ArgumentError("Applying functions of type $(nameof(typeof(f))) to batches of variates requires MeasureBase.mspace_ndims to be declared for the measure"))
 end
 
 massof(m::PushforwardMeasure) = massof(m.origin)
@@ -252,14 +288,14 @@ massof(m::PushforwardMeasure) = massof(m.origin)
 rand_impl(ctx::GenContext, ν::PushforwardMeasure) = ν.f(rand_impl(ctx, ν.origin))
 
 # Batches of pushforwards apply the function to the variates of a batch of
-# the origin, elementwise for scalar variates:
+# the origin:
 function batched_rand_impl(ctx::GenContext, ν::PushforwardMeasure, sz::Dims)
-    _pushfwd_batched_rand(ctx, ν, sz, mspace_flatsize(ν.origin))
+    _pushfwd_batched_rand(ctx, ν, sz, _static_ndims(ν.origin))
 end
-@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, ::Tuple{})
-    broadcast(ν.f, batched_rand_impl(ctx, ν.origin, sz))
+@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, k::StaticInteger)
+    _apply_batched(ν.f, batched_rand_impl(ctx, ν.origin, sz), k)
 end
-@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, ::Any)
+@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, ::NoMSpaceElementSize)
     _batched_rand_pointwise(ctx, ν, sz)
 end
 
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index d3969bcb..7fffad77 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -85,13 +85,13 @@ insupport(μ::WeightedMeasure, x) = insupport(μ.base, x)
 @inline transport_from_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, z::AbstractVector) where {S<:StdMeasure} =
     transport_from_std_with_rest(S, basemeasure(μ), z)
 
-@inline batched_transport_to_std(::Type{S}, μ::AbstractWeightedMeasure, X::AbstractArray) where {S<:StdMeasure} =
+@inline batched_transport_to_std(::Type{S}, μ::AbstractWeightedMeasure, X) where {S<:StdMeasure} =
     batched_transport_to_std(S, basemeasure(μ), X)
 @inline batched_transport_from_std(::Type{S}, μ::AbstractWeightedMeasure, Z::AbstractArray) where {S<:StdMeasure} =
     batched_transport_from_std(S, basemeasure(μ), Z)
-@inline batched_transport_to_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, X::AbstractArray) where {S<:StdMeasure} =
-    batched_transport_to_std_with_rest(S, basemeasure(μ), X)
-@inline batched_transport_from_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, Z::AbstractArray) where {S<:StdMeasure} =
-    batched_transport_from_std_with_rest(S, basemeasure(μ), Z)
+@inline batched_transport_to_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, X::AbstractArray, sz::Dims) where {S<:StdMeasure} =
+    batched_transport_to_std_with_rest(S, basemeasure(μ), X, sz)
+@inline batched_transport_from_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, Z::AbstractArray, sz::Dims) where {S<:StdMeasure} =
+    batched_transport_from_std_with_rest(S, basemeasure(μ), Z, sz)
 
 Adapt.adapt_structure(to, μ::WeightedMeasure) = WeightedMeasure(μ.logweight, Adapt.adapt(to, μ.base))
diff --git a/src/density-batched.jl b/src/density-batched.jl
index e9c75499..bb28947f 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -293,7 +293,8 @@ end
     reshape(X_flat, (dims..., Base.tail(size(X_flat))...))
 end
 
-@inline function _batched_split(A::AbstractArray, n_rows::Integer)
+@inline function _batched_split(A::AbstractArray, n::IntegerLike)
+    n_rows = dynamic(n)
     stream_idxs = axes(A, 1)
     if length(stream_idxs) < n_rows
         throw(ArgumentError("Variate streams too short during batched evaluation"))
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 9fbff0a6..0bf69750 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -58,14 +58,19 @@ end
 @inline transport_from_std(::Type{S}, μ::Dirac, z::AbstractVector) where {S<:StdMeasure} = μ.x
 @inline transport_from_std_with_rest(::Type{S}, μ::Dirac, z::AbstractVector) where {S<:StdMeasure} = μ.x, z
 
+@inline batched_transport_to_std(::Type{S}, ::Dirac, ::Number) where {S<:StdMeasure} = SVector{0,Bool}()
 function batched_transport_to_std(::Type{S}, μ::Dirac, X::AbstractArray) where {S<:StdMeasure}
     n = length(_value_flatsize(μ.x))
     similar(X, Bool, (0, ntuple(i -> size(X, n + i), Val(ndims(X) - n))...))
 end
 
 function batched_transport_from_std(::Type{S}, μ::Dirac, Z::AbstractArray) where {S<:StdMeasure}
-    X = similar(Z, eltype(μ.x), (size(μ.x)..., Base.tail(size(Z))...))
-    X .= μ.x
+    _const_variates(μ.x, Z)
+end
+@inline _const_variates(x::Number, ::AbstractVector) = x
+function _const_variates(x, Z::AbstractArray)
+    X = similar(Z, eltype(x), (size(x)..., Base.tail(size(Z))...))
+    X .= x
     return X
 end
 
diff --git a/src/transport-batched.jl b/src/transport-batched.jl
index 3ef0e0ff..b901478a 100644
--- a/src/transport-batched.jl
+++ b/src/transport-batched.jl
@@ -1,132 +1,210 @@
-# Batched transport over flat batches of variates: the leading dimensions of
-# a batch are the variate dimensions (see `mspace_flatsize`), all further
-# dimensions are batch dimensions. Streams of standard variates are batches
-# `(dof, batch dims...)`, consumed along their first dimension.
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
 
-"""
-    MeasureBase.batched_transport_to_std(::Type{S}, μ, X::AbstractArray)
-
-Batched form of [`MeasureBase.transport_to_std`](@ref): transports the
-flat batch `X` of variates of `μ` to a batch `(getdof(μ), batch dims...)`
-of variates of the standard measure type `S`.
+# Batched-first transport over flat batches of variates `(variate dims...,
+# batch dims...)`, zero batch dims meaning a single variate. Streams of
+# standard variates are batches `(dof, batch dims...)`, consumed along their
+# first dimension. Kernels know the variate rank of their measure (see
+# `mspace_ndims`), structural measures implement them once in terms of the
+# kernels of their components.
 
-The default implementation broadcasts the point transport for measures
-with scalar variates and maps it over the variate slices of `X`
-otherwise, in a host loop.
+"""
+    MeasureBase.batched_transport_to_std(::Type{S}, μ, X)
+
+Transport the flat batch `X` of variates of `μ` to a batch `(getdof(μ),
+batch dims...)` of variates of the standard measure type `S`. `X` may be a
+single variate, the result is a vector then.
+
+The default implementation broadcasts the point transport
+[`MeasureBase.transport_to_std`](@ref) for measures with scalar variates
+and maps it over the variate slices of `X` (in a host loop) for measures
+with array variates of a declared number of dimensions (see
+[`MeasureBase.mspace_ndims`](@ref)). Measure types with array variates
+should implement `batched_transport_to_std` directly.
 """
 function batched_transport_to_std end
 
-function batched_transport_to_std(::Type{S}, μ, X::AbstractArray) where {S<:StdMeasure}
-    _batched_to_std(S, μ, X, mspace_flatsize(μ))
+@inline function batched_transport_to_std(::Type{S}, μ, X) where {S<:StdMeasure}
+    _batched_to_std(S, μ, X, _static_ndims(μ))
 end
 
-@inline function _batched_to_std(::Type{S}, μ, X::AbstractArray, ::Tuple{}) where {S}
+@inline function _batched_to_std(::Type{S}, μ, X, ::StaticInteger{0}) where {S}
+    _as_stdstream_batch(broadcast(Base.Fix1(_ToStd{S}(), μ), X))
+end
+@inline function _batched_to_std(::Type{S}, μ, X::AbstractArray, ::StaticInteger{0}) where {S}
     _as_stdstream_batch(broadcast(Base.Fix1(_ToStd{S}(), μ), X))
 end
+@inline function _batched_to_std(::Type{S}, μ, X::AbstractArray, ::StaticInteger{K}) where {S,K}
+    _to_std_slices(S, μ, X, Val(K))
+end
+@noinline function _batched_to_std(::Type{S}, μ, X, ::NoMSpaceElementSize) where {S}
+    throw(ArgumentError("Batched transport requires MeasureBase.mspace_ndims to be declared for measures of type $(nameof(typeof(μ))) or MeasureBase.batched_transport_to_std to be implemented"))
+end
 
-function _batched_to_std(::Type{S}, μ, X::AbstractArray, sz::SizeLike) where {S}
-    _check_flatsize(X, sz)
-    stacked(map(Base.Fix1(_ToStd{S}(), μ), sliced(X, Val(length(sz)))))
+@inline _to_std_slices(::Type{S}, μ, X::AbstractArray{<:Any,K}, ::Val{K}) where {S,K} = _as_stdstream(transport_to_std(S, μ, X))
+@inline function _to_std_slices(::Type{S}, μ, X::AbstractArray, ::Val{K}) where {S,K}
+    stacked(map(Base.Fix1(_ToStd{S}(), μ), sliced(X, Val(K))))
 end
 
-function _batched_to_std(::Type{S}, μ, ::AbstractArray, ::NoMSpaceElementSize) where {S}
-    throw(ArgumentError("Batched transport requires measures of type $(nameof(typeof(μ))) to have a known variate size"))
+# Standard variates of scalar-variate measures form the first dimension of
+# a batch of streams:
+@inline _as_stdstream_batch(Z::AbstractArray) = _reshape_batch(Z, (static(1), _batch_dims(Z)...))
+@inline _as_stdstream_batch(z::Number) = SVector(z)
+@inline _drop_stdstream_dim(Z::AbstractArray) = _reshape_batch(Z, Base.tail(_batch_dims(Z)))
+
+# Sizes as tuples of (maybe static) integers, and reshapes that keep static
+# arrays static:
+@inline _batch_dims(A::AbstractArray) = _size_dims(maybestatic_size(A))
+@inline _reshape_batch(A::AbstractArray, dims::Tuple) = reshape(A, map(dynamic, dims))
+@inline _reshape_batch(A::StaticArray, dims::Tuple{Vararg{StaticInteger}}) = maybestatic_reshape(A, dims)
+
+# Merge the leading `N` dimensions of an array into one, `N == 0` adds a
+# leading dimension of size one:
+@inline function _merge_leading_dims(A::AbstractArray, ::StaticInteger{N}) where {N}
+    ndims(A) >= N || _throw_size_mismatch()
+    dims = _batch_dims(A)
+    lead = ntuple(i -> dims[i], Val(N))
+    _reshape_batch(A, (prod(lead), ntuple(i -> dims[N + i], Val(length(dims) - N))...))
 end
+@inline _merge_leading_dims(A::AbstractArray, ::StaticInteger{0}) = _reshape_batch(A, (static(1), _batch_dims(A)...))
 
-# Standard variates of scalar-variate measures form the first dimension:
-@inline _as_stdstream_batch(Z::AbstractArray) = reshape(Z, (1, size(Z)...))
-@inline _drop_stdstream_dim(Z::AbstractArray) = reshape(Z, Base.tail(size(Z)))
+# A flat batch of variates as a batch of streams, the variate dimensions
+# merged into the first dimension:
+@inline _as_stream_batch(X::AbstractArray, ::StaticInteger{K}) where {K} = _merge_leading_dims(X, static(K))
+@inline _as_stream_batch(x::Number, ::StaticInteger{0}) = SVector(x)
+@noinline function _as_stream_batch(X, ::NoMSpaceElementSize)
+    throw(ArgumentError("Concatenating batches of variates requires MeasureBase.mspace_ndims to be declared for the measures involved"))
+end
+
+# The standard variates of a single variate must form a vector:
+@inline _single_std(z::AbstractVector) = z
+@noinline function _single_std(z)
+    throw(ArgumentError("Transport of a single variate resulted in a batch of standard variates, the variate doesn't fit the measure"))
+end
+
+@noinline function _throw_std_length_mismatch()
+    throw(ArgumentError("Length of standard variates doesn't match the degrees of freedom of the measure"))
+end
 
 
 """
     MeasureBase.batched_transport_from_std(::Type{S}, μ, Z::AbstractArray)
 
-Batched form of [`MeasureBase.transport_from_std`](@ref): transports the
-batch `Z` of variates of the standard measure type `S`, of size
-`(getdof(μ), batch dims...)`, to a flat batch of variates of `μ`.
+Transport the batch `Z` of variates of the standard measure type `S`, of
+size `(getdof(μ), batch dims...)`, to a flat batch of variates of `μ`. A
+single stream `Z` yields a single variate.
 
-The default implementation broadcasts the point transport for measures
-with scalar variates and maps it over the columns of `Z` otherwise, in a
-host loop.
+The default implementation broadcasts the point transport
+[`MeasureBase.transport_from_std`](@ref) for measures with scalar
+variates and maps it over the columns of `Z` (in a host loop) for measures
+with array variates of a declared number of dimensions.
 """
 function batched_transport_from_std end
 
-function batched_transport_from_std(::Type{S}, μ, Z::AbstractArray) where {S<:StdMeasure}
-    _batched_from_std(S, μ, Z, mspace_flatsize(μ))
+@inline function batched_transport_from_std(::Type{S}, μ, Z::AbstractArray) where {S<:StdMeasure}
+    _batched_from_std(S, μ, Z, _static_ndims(μ))
 end
 
-@inline function _batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Tuple{}) where {S}
-    broadcast(Base.Fix1(_FromStd{S}(), μ), _drop_stdstream_dim(Z))
+@inline _batched_from_std(::Type{S}, μ, Z::AbstractArray, ::StaticInteger{0}) where {S} = _from_std_scalar(S, μ, Z)
+@inline _batched_from_std(::Type{S}, μ, Z::AbstractArray, ::StaticInteger{K}) where {S,K} = _from_std_columns(S, μ, Z)
+@noinline function _batched_from_std(::Type{S}, μ, ::AbstractArray, ::NoMSpaceElementSize) where {S}
+    throw(ArgumentError("Batched transport requires MeasureBase.mspace_ndims to be declared for measures of type $(nameof(typeof(μ))) or MeasureBase.batched_transport_from_std to be implemented"))
 end
 
-function _batched_from_std(::Type{S}, μ, Z::AbstractArray, sz::SizeLike) where {S}
-    xs = map(Base.Fix1(_FromStd{S}(), μ), sliced(Z, Val(1)))
-    reshape(stacked(xs), (map(dynamic, _size_dims(sz))..., Base.tail(size(Z))...))
+@inline _from_std_scalar(::Type{S}, μ, z::AbstractVector) where {S} = transport_from_std(S, μ, z[begin])
+@inline function _from_std_scalar(::Type{S}, μ, Z::AbstractArray) where {S}
+    broadcast(Base.Fix1(_FromStd{S}(), μ), _drop_stdstream_dim(Z))
 end
-
-function _batched_from_std(::Type{S}, μ, ::AbstractArray, ::NoMSpaceElementSize) where {S}
-    throw(ArgumentError("Batched transport requires measures of type $(nameof(typeof(μ))) to have a known variate size"))
+@inline _from_std_columns(::Type{S}, μ, z::AbstractVector) where {S} = transport_from_std(S, μ, z)
+@inline function _from_std_columns(::Type{S}, μ, Z::AbstractArray) where {S}
+    stacked(map(Base.Fix1(_FromStd{S}(), μ), sliced(Z, Val(1))))
 end
 
 
 """
-    MeasureBase.batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray)
-
-Batched form of [`MeasureBase.transport_to_std_with_rest`](@ref) for a
-batch `X` of flat vector streams (first dimension along the streams).
-
-Returns a tuple `(Z, X_μ, X_rest)` of the batch of standard variates, the
-batch of variates of `μ` consumed from the streams and the unconsumed rest
-of the streams.
+    MeasureBase.batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::Dims)
+
+Consume variates of `μ` from the batch `X` of flat vector streams (first
+dimension along the streams, further dimensions are batch dimensions), a
+batch of variates of size `sz` per stream, and transport them to the
+standard measure type `S`.
+
+Returns a tuple `(Z, X_rest)` of the standard variates as a batch
+`(getdof(μ) * prod(sz), batch dims...)` and the unconsumed rest of the
+streams. The default implementation consumes variates of the size given by
+[`MeasureBase.mspace_flatsize`](@ref) or
+[`MeasureBase.some_mspace_elsize`](@ref), a single stream with `sz == ()`
+goes through [`MeasureBase.transport_to_std_with_rest`](@ref). Measures
+whose variates are composed of the variates of other measures implement
+`batched_transport_to_std_with_rest` in terms of their components.
 """
 function batched_transport_to_std_with_rest end
 
-function batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray) where {S<:StdMeasure}
+function batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::Dims) where {S<:StdMeasure}
+    _to_std_with_rest_default(S, μ, X, sz)
+end
+
+function _to_std_with_rest_default(::Type{S}, μ, x::AbstractVector, ::Tuple{}) where {S}
+    z, _, x_rest = transport_to_std_with_rest(S, μ, x)
+    return z, x_rest
+end
+function _to_std_with_rest_default(::Type{S}, μ, X::AbstractArray, sz::Dims) where {S}
     vsz = _stream_consume_size(μ)
-    X_μ, X_rest = _batched_consume(X, vsz, ())
-    X_v = _consumed_variates(X_μ, vsz)
-    return batched_transport_to_std(S, μ, X_v), X_v, X_rest
+    X_μ, X_rest = _batched_consume(X, vsz, sz)
+    Z = batched_transport_to_std(S, μ, _consumed_variates(X_μ, vsz))
+    return _merge_multiplicity(Z, sz), X_rest
 end
 @inline _consumed_variates(X_μ::AbstractArray, ::Tuple{}) = _drop_stdstream_dim(X_μ)
 @inline _consumed_variates(X_μ::AbstractArray, ::SizeLike) = X_μ
 
+# Standard variates of `prod(sz)` variates per stream, `(dof, sz..., batch
+# dims...)`, as one stream chunk `(dof * prod(sz), batch dims...)`, and
+# back:
+@inline _merge_multiplicity(Z::AbstractArray, sz::Dims) = _merge_leading_dims(Z, static(1) + static(length(sz)))
+@inline _split_multiplicity(Z::AbstractArray, ::Tuple{}, n) = Z
+@inline function _split_multiplicity(Z::AbstractArray, sz::Dims, n)
+    _reshape_batch(Z, (n, sz..., Base.tail(_batch_dims(Z))...))
+end
 
-"""
-    MeasureBase.batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray)
-
-Batched form of [`MeasureBase.transport_from_std_with_rest`](@ref) for a
-batch `Z` of streams of standard variates (first dimension along the
-streams).
 
-Returns a tuple `(X, Z_rest)` of the flat batch of variates of `μ` and the
-unconsumed rest of the streams.
+"""
+    MeasureBase.batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::Dims)
+
+Consume standard variates of type `S` for a batch of variates of size
+`sz` per stream from the batch `Z` of streams of standard variates (first
+dimension along the streams) and transport them to `μ`.
+
+Returns a tuple `(X, Z_rest)` of the flat batch `(flat variate dims...,
+sz..., batch dims...)` of variates of `μ` and the unconsumed rest of the
+streams. The default implementation consumes [`getdof(μ)`](@ref) entries
+per variate, a single stream with `sz == ()` goes through
+[`MeasureBase.transport_from_std_with_rest`](@ref). Measures whose
+variates are composed of the variates of other measures implement
+`batched_transport_from_std_with_rest` in terms of their components.
 """
 function batched_transport_from_std_with_rest end
 
-function batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray) where {S<:StdMeasure}
-    _batched_from_std_with_rest_bydof(S, μ, Z, fast_dof(μ))
+function batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::Dims) where {S<:StdMeasure}
+    _from_std_with_rest_default(S, μ, Z, sz)
 end
 
-function _batched_from_std_with_rest_bydof(::Type{S}, μ, Z::AbstractArray, n::IntegerLike) where {S}
-    Z_μ, Z_rest = _batched_split(Z, dynamic(n))
-    return batched_transport_from_std(S, μ, Z_μ), Z_rest
+_from_std_with_rest_default(::Type{S}, μ, z::AbstractVector, ::Tuple{}) where {S} = transport_from_std_with_rest(S, μ, z)
+function _from_std_with_rest_default(::Type{S}, μ, Z::AbstractArray, sz::Dims) where {S}
+    _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ))
 end
 
-function _batched_from_std_with_rest_bydof(::Type{S}, μ, ::AbstractArray, ::AbstractNoDOF) where {S}
-    throw(ArgumentError("Batched transport from standard measures requires measures of type $(nameof(typeof(μ))) to implement MeasureBase.batched_transport_from_std_with_rest"))
+function _batched_from_std_bydof(::Type{S}, μ, Z::AbstractArray, sz::Dims, n::IntegerLike) where {S}
+    Z_μ, Z_rest = _batched_split(Z, _chunk_rows(n, sz))
+    return batched_transport_from_std(S, μ, _split_multiplicity(Z_μ, sz, n)), Z_rest
 end
-
-# A flat batch of variates as a batch of streams, the variate dimensions
-# merged into the first dimension:
-@inline function _as_stream_batch(X::AbstractArray, sz::SizeLike)
-    n = length(sz)
-    batch_dims = ntuple(i -> size(X, n + i), Val(ndims(X) - n))
-    reshape(X, (dynamic(size2length(sz)), batch_dims...))
+@noinline function _batched_from_std_bydof(::Type{S}, μ, ::AbstractArray, ::Dims, ::AbstractNoDOF) where {S}
+    throw(ArgumentError("Batched transport from standard measures requires measures of type $(nameof(typeof(μ))) to have fast degrees of freedom or to implement MeasureBase.batched_transport_from_std_with_rest"))
 end
+@inline _chunk_rows(n::IntegerLike, ::Tuple{}) = n
+@inline _chunk_rows(n::IntegerLike, sz::Dims) = dynamic(n) * prod(sz)
 
 
 """
-    MeasureBase.batched_transport_def(ν, μ, X::AbstractArray)
+    MeasureBase.batched_transport_def(ν, μ, X)
 
 Transport the flat batch `X` of variates of `μ` to a flat batch of
 variates of `ν`, via the standard measure type the preferences of `ν` and
@@ -135,10 +213,10 @@ batched transport.
 """
 function batched_transport_def end
 
-function batched_transport_def(ν, μ, X::AbstractArray)
+function batched_transport_def(ν, μ, X)
     S = _transport_pivot(ν, μ)
     Z = batched_transport_to_std(S, μ, X)
-    Y, Z_rest = batched_transport_from_std_with_rest(S, ν, Z)
+    Y, Z_rest = batched_transport_from_std_with_rest(S, ν, Z, ())
     if size(Z_rest, 1) != 0
         throw(ArgumentError("Degrees of freedom of source and target measure of a transport don't match"))
     end
@@ -151,7 +229,7 @@ end
 # whole. Fused broadcast arguments are materialized first, static arrays
 # are transported point by point.
 function Broadcast.broadcasted(f::TransportFunction, X::AbstractArray)
-    _broadcast_transport(f, X, _flat_storage(X), mspace_flatsize(f.μ), mspace_flatsize(f.ν))
+    _broadcast_transport(f, X, _flat_storage(X), _static_ndims(f.μ), _static_ndims(f.ν))
 end
 
 function Broadcast.broadcasted(f::TransportFunction, bc::Broadcast.Broadcasted)
@@ -160,10 +238,9 @@ end
 
 Broadcast.broadcasted(f::TransportFunction, X::StaticArray) = map(_Pointwise(f), X)
 
-function _broadcast_transport(f::TransportFunction, X, X_flat::AbstractArray, sz_μ::SizeLike, sz_ν::SizeLike)
-    _check_flatsize(X_flat, sz_μ)
+function _broadcast_transport(f::TransportFunction, X, X_flat::AbstractArray, ::StaticInteger, ::StaticInteger{K}) where {K}
     Y_flat = batched_transport_def(f.ν, f.μ, X_flat)
-    return _batch_variates(Y_flat, f.ν)
+    return _batch_variates(Y_flat, f.ν, Val(K))
 end
 
 _broadcast_transport(f::TransportFunction, X, ::Any, ::Any, ::Any) = map(_Pointwise(f), X)
@@ -177,5 +254,10 @@ end
 
 # The batch of variates in the layout of the target measure over the flat
 # result, nested powers included:
-@inline _batch_variates(Y::AbstractArray, ν) = _nest_leaf(Y, mspace_flatsize(ν))
-@inline _batch_variates(Y::AbstractArray, ν::PowerMeasure) = sliced(_pwr_variate(ν, Y), Val(length(pwr_axes(ν))))
+@inline _batch_variates(Y::AbstractArray, ν, ::Val{K}) where {K} = _nest_batch(Y, Val(K))
+@inline function _batch_variates(Y::AbstractArray, ν::PowerMeasure, ::Val)
+    sliced(_pwr_variate(ν, Y), Val(length(pwr_axes(ν))))
+end
+@inline _nest_batch(Y::AbstractArray, ::Val{0}) = Y
+@inline _nest_batch(Y::AbstractArray{<:Any,K}, ::Val{K}) where {K} = Y
+@inline _nest_batch(Y::AbstractArray, ::Val{K}) where {K} = sliced(Y, Val(K))
diff --git a/src/transport.jl b/src/transport.jl
index af9116b7..9fde5ff0 100644
--- a/src/transport.jl
+++ b/src/transport.jl
@@ -13,9 +13,9 @@ to `rand(ν)`. `f` supports `InverseFunctions.inverse` and
 Measures are transported via standard measures: `x` is transported to the
 standard measure type that the preferences of `ν` and `μ` promote to (see
 [`MeasureBase.preferred_stdmeasure`](@ref)) and from there to `ν`.
-Broadcasting `f` over an array of variates with flat storage (see
-[`MeasureBase.mspace_flatsize`](@ref)), or over the flat storage of a
-batch of variates, transports the whole batch at once. A standard measure
+Broadcasting `f` over an array of variates with flat storage, or over
+the flat storage of a batch of variates (see
+[`MeasureBase.mspace_ndims`](@ref)), transports the whole batch at once. A standard measure
 type like `StdUniform` or `StdNormal` may also be used directly as the
 source or target:
 
@@ -33,12 +33,16 @@ the measure has degrees of freedom otherwise.
 To support transport for a measure type, specialize
 [`MeasureBase.transport_to_std`](@ref) and
 [`MeasureBase.transport_from_std`](@ref) for its preferred standard measure
-type. Measures whose variates are composed of the variates of other
-measures specialize the stream forms
+type, and declare [`MeasureBase.mspace_ndims`](@ref) for array variates.
+Measure types with array variates should also implement the batched forms
+[`MeasureBase.batched_transport_to_std`](@ref) and
+[`MeasureBase.batched_transport_from_std`](@ref), which transport whole
+batches of variates. Measures whose variates are composed of the variates
+of other measures specialize the stream forms
 [`MeasureBase.transport_to_std_with_rest`](@ref) and
-[`MeasureBase.transport_from_std_with_rest`](@ref) instead.
-[`MeasureBase.transport_def`](@ref) may be specialized for pairs of
-measure types with a direct transport.
+[`MeasureBase.transport_from_std_with_rest`](@ref) instead (and their
+batched forms). [`MeasureBase.transport_def`](@ref) may be specialized
+for pairs of measure types with a direct transport.
 """
 function transport_to end
 export transport_to
@@ -280,8 +284,8 @@ function _from_std_with_rest_bydof(::Type{S}, μ, z::AbstractVector, ::AbstractN
 end
 
 # Scalar-variate measures take their standard variate as a number:
-@inline _chunk_as_variate(μ, z) = _chunk_as_variate(z, mspace_flatsize(μ))
-@inline _chunk_as_variate(z::AbstractVector, ::Tuple{}) = z[begin]
+@inline _chunk_as_variate(μ, z) = _chunk_as_variate(z, _static_ndims(μ))
+@inline _chunk_as_variate(z::AbstractVector, ::StaticInteger{0}) = z[begin]
 @inline _chunk_as_variate(z::AbstractVector, ::Any) = z
 
 
@@ -315,10 +319,10 @@ end
 
 function _std_tp_partner(::Type{M}, μ) where {M<:StdMeasure}
     m = asmeasure(μ)
-    _std_tp_partner_bysize(M, mspace_flatsize(m), m)
+    _std_tp_partner_byrank(M, _static_ndims(m), m)
 end
-_std_tp_partner_bysize(::Type{M}, ::Tuple{}, μ) where {M<:StdMeasure} = M()
-_std_tp_partner_bysize(::Type{M}, ::Any, μ) where {M<:StdMeasure} = M()^some_dof(μ)
+_std_tp_partner_byrank(::Type{M}, ::StaticInteger{0}, μ) where {M<:StdMeasure} = M()
+_std_tp_partner_byrank(::Type{M}, ::Any, μ) where {M<:StdMeasure} = M()^some_dof(μ)
 
 
 # Element-wise transport kernels for broadcasts and maps:
diff --git a/test/runtests.jl b/test/runtests.jl
index a6cf4e8c..5a427976 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -19,6 +19,7 @@ include("shape_contract.jl")
 include("logdensities.jl")
 include("numtype.jl")
 include("transport.jl")
+include("transport_batched.jl")
 include("smf.jl")
 include("domains.jl")
 
diff --git a/test/transport_batched.jl b/test/transport_batched.jl
new file mode 100644
index 00000000..db91ef5a
--- /dev/null
+++ b/test/transport_batched.jl
@@ -0,0 +1,171 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using MeasureBase
+using MeasureBase: StdUniform, StdExponential, StdLogistic, StdNormal, Dirac, GenContext
+using MeasureBase: productmeasure, pushfwd, mcombine, weightedmeasure, mbind, PushfwdRootMeasure
+using MeasureBase: transport_to_std, transport_from_std, transport_to_std_with_rest
+using MeasureBase: batched_transport_to_std, batched_transport_from_std
+using MeasureBase: batched_transport_to_std_with_rest, batched_transport_from_std_with_rest
+using MeasureBase: batched_rand_impl
+using InverseFunctions: inverse
+using ArraysOfArrays: sliced, flatview
+using StaticArrays: SVector
+using Distributions: MvNormal, LogNormal, logpdf
+using AffineMaps: Mul, MulAdd
+using JLArrays
+
+@testset "batched transport" begin
+    stdn_to_u = transport_to(StdUniform(), StdNormal())
+
+    @testset "several variates per stream" begin
+        X = vcat(randn(2, 6), rand(3, 6))
+        Z, R = batched_transport_to_std_with_rest(StdUniform, StdNormal(), X, (2,))
+        @test size(Z) == (2, 6) && size(R) == (3, 6)
+        @test Z ≈ stdn_to_u.(X[1:2, :])
+        Xb, Rb = batched_transport_from_std_with_rest(StdUniform, StdNormal(), Z, (2,))
+        @test Xb ≈ X[1:2, :] && size(Rb) == (0, 6)
+
+        # Powers consume their base with their size as multiplicity:
+        Zp, Rp = batched_transport_to_std_with_rest(StdUniform, StdNormal()^2, X, ())
+        @test Zp ≈ Z && size(Rp) == (3, 6)
+        X2 = vcat(X, X)
+        Zp2, Rp2 = batched_transport_to_std_with_rest(StdUniform, StdNormal()^2, X2, (2,))
+        @test size(Zp2) == (4, 6) && size(Rp2) == (6, 6)
+        @test Zp2 ≈ stdn_to_u.(X2[1:4, :])
+        Xp2, _ = batched_transport_from_std_with_rest(StdUniform, StdNormal()^2, Zp2, (2,))
+        @test Xp2 ≈ reshape(X2[1:4, :], (2, 2, 6))
+
+        # Combined measures split the rows of each variate by component:
+        m = mcombine(vcat, StdNormal()^2, StdUniform()^3)
+        Z1 = batched_transport_to_std(StdUniform, m, X)
+        Zm, Rm = batched_transport_to_std_with_rest(StdUniform, m, X2, (2,))
+        @test size(Rm) == (0, 6) && Zm ≈ vcat(Z1, Z1)
+        Xm, _ = batched_transport_from_std_with_rest(StdUniform, m, Zm, (2,))
+        @test Xm ≈ reshape(X2, (5, 2, 6))
+
+        # Tuple products consume several variates via their degrees of freedom:
+        Pt = productmeasure((StdNormal(), StdExponential()^2))
+        Zt = rand(6, 4)
+        Xt, Rt = batched_transport_from_std_with_rest(StdUniform, Pt, Zt, (2,))
+        @test size(Xt[1]) == (2, 4) && size(Xt[2]) == (2, 2, 4) && size(Rt) == (0, 4)
+        for j in 1:4, i in 1:2
+            a, b = transport_from_std(StdUniform, Pt, Zt[(3i - 2):(3i), j])
+            @test_throws ArgumentError transport_from_std(StdUniform, Pt, Zt[:, j])
+            @test Xt[1][i, j] ≈ a && Xt[2][:, i, j] ≈ b
+        end
+
+        @test_throws ArgumentError batched_transport_to_std_with_rest(StdUniform, StdNormal()^2, X, (3,))
+    end
+
+    @testset "tuple and named tuple products" begin
+        Pt = productmeasure((StdNormal(), StdExponential()^2))
+        Xt = (randn(4), rand(2, 4))
+        Zt = batched_transport_to_std(StdUniform, Pt, Xt)
+        @test size(Zt) == (3, 4)
+        @test Zt ≈ stack([transport_to_std(StdUniform, Pt, (Xt[1][j], Xt[2][:, j])) for j in 1:4])
+        Xr = batched_transport_from_std(StdUniform, Pt, Zt)
+        @test Xr[1] ≈ Xt[1] && Xr[2] ≈ Xt[2]
+        Pn = productmeasure((a = StdNormal(), b = StdExponential()^2))
+        Zn = batched_transport_to_std(StdUniform, Pn, (a = Xt[1], b = Xt[2]))
+        @test Zn ≈ Zt
+        Xn = batched_transport_from_std(StdUniform, Pn, Zn)
+        @test Xn.a ≈ Xt[1] && Xn.b ≈ Xt[2]
+        @test batched_transport_to_std(StdUniform, Pt, (Xt[1][1], Xt[2][:, 1])) ≈ Zt[:, 1]
+        x1 = batched_transport_from_std(StdUniform, Pt, Zt[:, 1])
+        @test x1[1] ≈ Xt[1][1] && x1[2] ≈ Xt[2][:, 1]
+        @test_throws ArgumentError batched_transport_from_std(StdUniform, Pt, rand(4, 4))
+    end
+
+    @testset "array products of array-variate marginals" begin
+        P = productmeasure([weightedmeasure(log(i), StdNormal()^2) for i in 1:3])
+        X = randn(2, 3, 5)
+        Z = batched_transport_to_std(StdUniform, P, X)
+        f = transport_to(StdUniform()^6, P)
+        @test size(Z) == (6, 5) && Z ≈ stack([f(X[:, :, j]) for j in 1:5])
+        @test batched_transport_from_std(StdUniform, P, Z) ≈ X
+        Y = f.(sliced(X, Val(2)))
+        @test flatview(Y) ≈ Z
+        @test flatview(inverse(f).(Y)) ≈ X
+        xs = [randn(2) for _ in 1:3]
+        z = transport_to_std(StdUniform, P, xs)
+        @test z ≈ f(stack(xs))
+        xn = transport_from_std(StdUniform, P, z)
+        @test length(xn) == 3 && all(xn[i] ≈ xs[i] for i in 1:3)
+        @test_throws ArgumentError batched_transport_to_std(StdUniform, P, randn(2, 2, 5))
+        @test_throws ArgumentError batched_transport_from_std(StdUniform, P, rand(5, 5))
+    end
+
+    @testset "streams with value-dependent sizes" begin
+        f_β(a) = StdNormal()^length(a)
+        μb = mbind(f_β, StdUniform()^1, vcat)
+        m = mcombine(vcat, μb, StdExponential())
+        X = vcat(rand(1, 4), randn(1, 4), rand(1, 4))
+        Z = batched_transport_to_std(StdUniform, m, X)
+        @test size(Z) == (3, 4)
+        @test Z ≈ stack([transport_to_std(StdUniform, m, X[:, j]) for j in 1:4])
+        @test batched_transport_from_std(StdUniform, m, Z) ≈ X
+        P = μb^2
+        x = vcat(rand(1), randn(1), rand(1), randn(1), rand(2))
+        z, x_μ, x_rest = transport_to_std_with_rest(StdUniform, P, x)
+        @test length(z) == 4 && length(x_μ) == 4 && length(x_rest) == 2
+        @test z ≈ transport_to_std(StdUniform, P, [x[1:2], x[3:4]])
+    end
+
+    @testset "elementwise pushforwards" begin
+        νe = pushfwd(Base.BroadcastFunction(exp), StdNormal()^3)
+        @test MeasureBase.mspace_ndims(typeof(νe)) == 1
+        Ye = exp.(randn(3, 4))
+        @test logdensities(νe, Ye) ≈ [logdensityof(νe, Ye[:, j]) for j in 1:4]
+        @test logdensityof(νe, Ye[:, 1]) ≈ sum(logpdf.(LogNormal(), Ye[:, 1]))
+        fe = transport_to(StdUniform()^3, νe)
+        @test flatview(fe.(sliced(Ye, Val(1)))) ≈ stack(map(fe, eachcol(Ye)))
+        @test flatview(inverse(fe).(fe.(sliced(Ye, Val(1))))) ≈ Ye
+        @test size(batched_rand_impl(GenContext{Float64}(), νe, (5,))) == (3, 5)
+        νr = pushfwd(Base.BroadcastFunction(exp), StdNormal()^3, PushfwdRootMeasure())
+        @test logdensities(νr, Ye) ≈ [logdensityof(νr, Ye[:, j]) for j in 1:4]
+        JLArrays.allowscalar(false)
+        @test Array(logdensities(νe, JLArray(Ye))) ≈ logdensities(νe, Ye)
+        @test Array(flatview(fe.(sliced(JLArray(Ye), Val(1))))) ≈ flatview(fe.(sliced(Ye, Val(1))))
+    end
+
+    @testset "affine pushforwards" begin
+        A = [2.0 0.5; 0.0 1.5]
+        b = [1.0, -1.0]
+        ν = pushfwd(MulAdd(A, b), StdNormal()^2)
+        @test MeasureBase.mspace_ndims(typeof(ν)) == 1
+        Y = randn(2, 5)
+        @test logdensities(ν, Y) ≈ [logpdf(MvNormal(b, A * A'), Y[:, j]) for j in 1:5]
+        @test logdensityof(ν, Y[:, 1]) ≈ logpdf(MvNormal(b, A * A'), Y[:, 1])
+        f = transport_to(StdUniform()^2, ν)
+        @test flatview(f.(sliced(Y, Val(1)))) ≈ stack(map(f, eachcol(Y)))
+        @test flatview(inverse(f).(f.(sliced(Y, Val(1))))) ≈ Y
+        @test size(batched_rand_impl(GenContext{Float64}(), ν, (7,))) == (2, 7)
+        νs = pushfwd(Mul(2.0), StdNormal())
+        @test logdensities(νs, Y[1, :]) ≈ logdensityof.(Ref(νs), Y[1, :])
+    end
+
+    @testset "single variates through batched forms" begin
+        approx(a::Tuple, b::Tuple) = all(map(approx, a, b))
+        approx(a, b) = a ≈ b
+        for μ in (
+            StdNormal(),
+            StdNormal()^3,
+            productmeasure((StdNormal(), StdExponential()^2)),
+            mcombine(vcat, StdNormal()^2, StdUniform()^3),
+            weightedmeasure(0.3, StdNormal()^2),
+            pushfwd(Base.BroadcastFunction(exp), StdNormal()^2),
+        )
+            x = rand(μ)
+            z = MeasureBase._as_stdstream(transport_to_std(StdUniform, μ, x))
+            zb = batched_transport_to_std(StdUniform, μ, x)
+            @test zb isa AbstractVector && zb ≈ z
+            @test approx(batched_transport_from_std(StdUniform, μ, z), x)
+            @test approx(transport_from_std(StdUniform, μ, MeasureBase._chunk_as_variate(μ, z)), x)
+        end
+        @test batched_transport_to_std(StdUniform, Dirac(1.0), 1.0) == SVector{0,Bool}()
+        @test batched_transport_from_std(StdUniform, Dirac(2.0), SVector{0,Bool}()) == 2.0
+        @test_throws ArgumentError transport_to_std(StdUniform, StdNormal()^3, randn(3, 2))
+    end
+end

From e689389c483d06ab70d2db117efbfee767604638 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 12:15:50 +0200
Subject: [PATCH 095/122] Generate random variates batched-first

batched_rand_impl is now the primary extension point, a single variate is
a batch with zero batch dimensions and rand_impl derives from it. Defaults
draw standard variates and transport them in flat batches, or generate
variates one by one for measures without a standard transport, the
markers on the default paths keep the two extension points from
recursing into each other. Powers generate through their base measure
with their size as additional batch dimensions, batches of tuple and
named tuple products are tuples of marginal batches, and variates of
powers of such products are struct arrays. Superpositions and spike
mixtures mask by variate rank. Adds test/cuda, an opt-in CUDA test
runner for densities, transports and random variates on the device.

Created by generative AI.
---
 .../distribution_measure.jl                   |  12 +-
 src/combinators/combined.jl                   |   7 +
 src/combinators/power.jl                      |  48 +++----
 src/combinators/product.jl                    |   6 +
 src/combinators/spikemixture.jl               |   6 +-
 src/combinators/superpose.jl                  |   6 +-
 src/rand.jl                                   | 132 +++++++++---------
 src/standard/stdmeasure.jl                    |   1 +
 test/cuda/Project.toml                        |  13 ++
 test/cuda/runtests.jl                         | 116 +++++++++++++++
 test/rand_batched.jl                          | 102 ++++++++++++++
 test/runtests.jl                              |   1 +
 12 files changed, 346 insertions(+), 104 deletions(-)
 create mode 100644 test/cuda/Project.toml
 create mode 100644 test/cuda/runtests.jl
 create mode 100644 test/rand_batched.jl

diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index 9bc59610..d2ecc0bf 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -21,19 +21,23 @@ MeasureBase.rand_impl(ctx::GenContext, m::DistributionMeasure) =
 MeasureBase.batched_rand_impl(ctx::GenContext, m::DistributionMeasure, sz::Dims) =
     _flat_powrand(get_rng(ctx), get_precision(ctx), m.obj, sz)
 
-function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution{<:ArrayLikeVariate{0}}, sz::Dims) where {T<:Real}
+# A single variate for zero batch dimensions, flat batches otherwise:
+_flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution, ::Tuple{}) where {T<:Real} = convert_realtype(T, rand(rng, d))
+_flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution, sz::Dims) where {T<:Real} = _flat_powrand_batch(rng, T, d, sz)
+
+function _flat_powrand_batch(rng::AbstractRNG, ::Type{T}, d::Distribution{<:ArrayLikeVariate{0}}, sz::Dims) where {T<:Real}
     convert_realtype(T, reshape(rand(rng, d, prod(sz)), sz...))
 end
 
-function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution{<:ArrayLikeVariate{1}}, sz::Dims) where {T<:Real}
+function _flat_powrand_batch(rng::AbstractRNG, ::Type{T}, d::Distribution{<:ArrayLikeVariate{1}}, sz::Dims) where {T<:Real}
     convert_realtype(T, reshape(rand(rng, d, prod(sz)), size(d)..., sz...))
 end
 
-function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::ReshapedDistribution{N,<:Any,<:Distribution{<:ArrayLikeVariate{1}}}, sz::Dims) where {T<:Real,N}
+function _flat_powrand_batch(rng::AbstractRNG, ::Type{T}, d::ReshapedDistribution{N,<:Any,<:Distribution{<:ArrayLikeVariate{1}}}, sz::Dims) where {T<:Real,N}
     convert_realtype(T, reshape(rand(rng, d.dist, prod(sz)), d.dims..., sz...))
 end
 
-function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution, sz::Dims) where {T<:Real}
+function _flat_powrand_batch(rng::AbstractRNG, ::Type{T}, d::Distribution, sz::Dims) where {T<:Real}
     flatview(ArrayOfSimilarArrays(convert_realtype(T, rand(rng, d, sz))))
 end
 
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index 21185aab..0187ba22 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -247,6 +247,13 @@ end
 
 rand_impl(ctx::GenContext, μ::CombinedMeasure) = μ.f_c(rand_impl(ctx, μ.α), rand_impl(ctx, μ.β))
 
+batched_rand_impl(ctx::GenContext, μ::CombinedMeasure, sz::Dims) = _batched_rand_pointwise(ctx, μ, sz)
+
+# Batches of merge-combined measures merge the named tuples of batches:
+function batched_rand_impl(ctx::GenContext, μ::CombinedMeasure{typeof(merge)}, sz::Dims)
+    merge(batched_rand_impl(ctx, μ.α, sz), batched_rand_impl(ctx, μ.β, sz))
+end
+
 # Batches of vcat-combined measures are concatenated along the streams:
 function batched_rand_impl(ctx::GenContext, μ::CombinedMeasure{typeof(vcat)}, sz::Dims)
     _combined_batched_rand(ctx, μ, sz, _static_ndims(μ.α), _static_ndims(μ.β))
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index d5e959b7..baaf033e 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -56,40 +56,23 @@ function _cartidxs(axs::Tuple{Vararg{AbstractUnitRange,N}}) where {N}
 end
 
 # Variates of powers are generated as one flat batch of variates of the
-# innermost base measure, in the layout of the flat variate storage:
+# base measure, with the power's size as additional batch dimensions. Base
+# measures without fixed variate sizes generate their variates one by one.
 
-rand_impl(ctx::GenContext, μ::PowerMeasure) = _pwr_rand(ctx, μ, mspace_flatsize(μ))
-
-function _pwr_rand(ctx::GenContext, μ::PowerMeasure, sz_flat::SizeLike)
-    ν, _ = _pwr_unwrap(μ)
-    _pwr_variate(μ, batched_rand_impl(ctx, ν, _pwr_batch_dims(sz_flat, mspace_flatsize(ν))))
-end
-
-function _pwr_rand(ctx::GenContext, μ::PowerMeasure, ::NoMSpaceElementSize)
+rand_impl(ctx::GenContext, μ::PowerMeasure) = _pwr_rand(ctx, μ, fixed_stream_size(pwr_base(μ)))
+_pwr_rand(ctx::GenContext, μ::PowerMeasure, ::True) = _pwr_variate(μ, batched_rand_impl(ctx, μ, ()))
+function _pwr_rand(ctx::GenContext, μ::PowerMeasure, ::False)
     ν = pwr_base(μ)
     map(_ -> rand_impl(ctx, ν), _cartidxs(pwr_axes(μ)))
 end
 
-# The power dimensions of a flat size, after the flat dimensions of the
-# innermost base measure:
-@inline function _pwr_batch_dims(sz_flat::SizeLike, sz_base::SizeLike)
-    dims = map(dynamic, _size_dims(sz_flat))
-    n = length(sz_base)
-    ntuple(i -> dims[n + i], Val(length(dims) - n))
-end
-
 function batched_rand_impl(ctx::GenContext, μ::PowerMeasure, sz::Dims)
-    _pwr_batched_rand(ctx, μ, sz, mspace_flatsize(μ))
+    _pwr_batched_rand(ctx, μ, sz, fixed_stream_size(pwr_base(μ)))
 end
-
-function _pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::Dims, sz_flat::SizeLike)
-    ν, _ = _pwr_unwrap(μ)
-    batched_rand_impl(ctx, ν, (_pwr_batch_dims(sz_flat, mspace_flatsize(ν))..., sz...))
-end
-
-function _pwr_batched_rand(::GenContext, μ::PowerMeasure, ::Dims, ::NoMSpaceElementSize)
-    throw(ArgumentError("Batched random variate generation for powers of measures of type $(nameof(typeof(pwr_base(μ)))) requires a known variate size"))
+function _pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::Dims, ::True)
+    batched_rand_impl(ctx, pwr_base(μ), (_dynamic_dims(pwr_size(μ))..., sz...))
 end
+_pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::Dims, ::False) = _batched_rand_pointwise(ctx, μ, sz)
 
 marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes)
 
@@ -343,9 +326,18 @@ function _pwr_from_std_with_rest(::Type{S}, μ, z, ::AbstractNoDOF) where {S}
     _marginals_from_std_with_rest(S, marginals(μ), z)
 end
 
-# The nested variate layout of a power over its flat storage:
-@inline _pwr_variate(μ::PowerMeasure, A::AbstractArray) = _pwr_nest(pwr_base(μ), _pwr_variate(pwr_base(μ), A))
+# The nested variate layout of a power over its flat storage, batches of
+# tuple variates become struct arrays:
+@inline _pwr_variate(μ::PowerMeasure, A::AbstractArray) = _pwr_variate_impl(μ, A)
+@inline _pwr_variate(μ::PowerMeasure, A::Union{Tuple,NamedTuple}) = _pwr_variate_impl(μ, A)
+@inline _pwr_variate_impl(μ::PowerMeasure, A) = _pwr_nest(pwr_base(μ), _pwr_variate(pwr_base(μ), A))
 @inline _pwr_variate(ν, A::AbstractArray) = _nest_leaf(A, _static_ndims(ν))
+@inline function _pwr_variate(ν::ProductMeasure{<:Tuple}, X::Tuple)
+    StructArray(map((m, Xi) -> _nest_leaf(Xi, _static_ndims(m)), marginals(ν), X))
+end
+@inline function _pwr_variate(ν::ProductMeasure{<:NamedTuple{names}}, X::NamedTuple{names}) where {names}
+    StructArray(NamedTuple{names}(map((m, Xi) -> _nest_leaf(Xi, _static_ndims(m)), values(marginals(ν)), values(X))))
+end
 @inline _nest_leaf(A::AbstractArray, ::StaticInteger{0}) = A
 @inline _nest_leaf(A::AbstractArray, ::NoMSpaceElementSize) = A
 @inline _nest_leaf(A::AbstractArray{<:Any,N}, ::StaticInteger{N}) where {N} = A
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 2adcb3b7..bed70a5b 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -67,6 +67,12 @@ end
 proxy(μ::ProductMeasure{<:FillArrays.Fill}) =
     powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ)))
 
+# Batches of tuple and named tuple products are tuples resp. named tuples
+# of marginal batches:
+function batched_rand_impl(ctx::GenContext, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, sz::Dims)
+    map(m -> batched_rand_impl(ctx, m, sz), marginals(μ))
+end
+
 # Batches of tuple and named tuple variates are tuples resp. named tuples
 # of batches, the marginal densities add up lazily:
 for (bhead, head) in [(:batched_logdensityof_impl, :logdensityof_impl), (:batched_logdensity_def, :logdensity_def)]
diff --git a/src/combinators/spikemixture.jl b/src/combinators/spikemixture.jl
index 58d8ad19..3f05b125 100644
--- a/src/combinators/spikemixture.jl
+++ b/src/combinators/spikemixture.jl
@@ -38,11 +38,11 @@ function rand_impl(ctx::GenContext, μ::SpikeMixture)
 end
 
 function batched_rand_impl(ctx::GenContext, μ::SpikeMixture, sz::Dims)
-    _spike_batched_rand(ctx, μ, sz, mspace_flatsize(μ.m))
+    _spike_batched_rand(ctx, μ, sz, _static_ndims(μ.m))
 end
-function _spike_batched_rand(ctx::GenContext, μ::SpikeMixture, sz::Dims, sz_flat::SizeLike)
+function _spike_batched_rand(ctx::GenContext, μ::SpikeMixture, sz::Dims, k::StaticInteger)
     X = batched_rand_impl(ctx, μ.m, sz)
-    return ifelse.(_batch_mask(_rand_bulk(ctx, sz) .< μ.w, sz_flat), X, zero(eltype(X)))
+    return ifelse.(_batch_mask(_rand_bulk(ctx, sz) .< μ.w, k), X, zero(eltype(X)))
 end
 _spike_batched_rand(ctx::GenContext, μ::SpikeMixture, sz::Dims, ::NoMSpaceElementSize) =
     _batched_rand_pointwise(ctx, μ, sz)
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index 2da33ce1..d0370eed 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -149,12 +149,12 @@ end
 # Batches of superpositions draw a batch from each component and select
 # by mass, branch-free:
 function batched_rand_impl(ctx::GenContext, μ::SuperpositionMeasure, sz::Dims)
-    _superpose_batched_rand(ctx, μ, sz, mspace_flatsize(μ))
+    _superpose_batched_rand(ctx, μ, sz, _static_ndims(μ))
 end
-function _superpose_batched_rand(ctx::GenContext, μ::SuperpositionMeasure, sz::Dims, sz_flat::SizeLike)
+function _superpose_batched_rand(ctx::GenContext, μ::SuperpositionMeasure, sz::Dims, k::StaticInteger)
     components = values(μ.components)
     masses, total = _component_masses(μ)
-    thresholds = _batch_mask(_rand_bulk(ctx, sz) .* total, sz_flat)
+    thresholds = _batch_mask(_rand_bulk(ctx, sz) .* total, k)
     X = batched_rand_impl(ctx, first(components), sz)
     csum = first(masses)
     for (mass, c) in Iterators.drop(zip(masses, components), 1)
diff --git a/src/rand.jl b/src/rand.jl
index 5168a087..a1b48063 100644
--- a/src/rand.jl
+++ b/src/rand.jl
@@ -1,7 +1,8 @@
 # Random variate generation is parameterized by a `GenContext` carrying the
 # random number generator, the numerical precision and the compute unit.
-# Batches of variates are generated in flat form on the compute unit, single
-# variates of powers are drawn as one batch and reshaped.
+# Variates are generated as flat batches `(variate dims..., batch dims...)`
+# on the compute unit, a single variate is a batch with zero batch
+# dimensions.
 
 """
     rand([rng::AbstractRNG], [T::Type{<:AbstractFloat}], μ::AbstractMeasure)
@@ -12,11 +13,11 @@ Generate a random variate of `μ`.
 The generative context `ctx` (see `HeterogeneousComputing.GenContext`)
 determines the random number generator, the numerical precision (`Float64`
 by default) and the compute unit that array-valued variates are generated
-on. The variates of powers of measures with a known flat variate size are
-generated in one batch.
+on. Variates of powers of measures are generated as one flat batch of
+variates of the base measure.
 
-Measure types should specialize [`MeasureBase.rand_impl`](@ref) and
-[`MeasureBase.batched_rand_impl`](@ref) instead of `rand`.
+Measure types should specialize [`MeasureBase.batched_rand_impl`](@ref)
+instead of `rand`.
 """
 Base.rand(ctx::GenContext, μ::AbstractMeasure) = rand_impl(ctx, μ)
 
@@ -28,83 +29,74 @@ Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractMeasure) where {T<:AbstractFl
 @inline Random.rand!(d::AbstractMeasure, args...) = rand!(Random.default_rng(), d, args...)
 
 
+"""
+    MeasureBase.batched_rand_impl(ctx::GenContext, μ, sz::Dims)
+
+Generate a batch of random variates of `μ` of batch size `sz` in flat
+form, an array `(variate dims..., sz...)`, or a single variate for
+`sz == ()`. Batches of tuple and named tuple variates are tuples resp.
+named tuples of batches.
+
+This is the primary extension point for random variate generation. The
+default implementation draws a batch of variates of the preferred
+standard measure of `μ` and transports it to `μ` (see
+[`MeasureBase.batched_transport_from_std`](@ref)), or generates the
+variates one by one via [`MeasureBase.rand_impl`](@ref) if `μ` has no
+standard transport.
+"""
+function batched_rand_impl end
+
 """
     MeasureBase.rand_impl(ctx::GenContext, μ)
 
 Generate one random variate of `μ` in the generative context `ctx`.
 
-The default implementation draws a variate of the preferred standard
-measure of `μ` and transports it to `μ`. Measure types with a more direct
-way of generating variates specialize `rand_impl`, and should specialize
-[`MeasureBase.batched_rand_impl`](@ref) as well where batches can be
-generated in a more direct way, too.
+The default implementation generates a batch with zero batch dimensions
+via [`MeasureBase.batched_rand_impl`](@ref). Measure types with a more
+direct way of generating single variates may specialize `rand_impl`.
 """
 function rand_impl end
 
-function rand_impl(ctx::GenContext, μ)
-    _rand_via_std(ctx, μ, preferred_stdmeasure(μ), fast_dof(μ), mspace_flatsize(μ))
-end
+# The marker tells the defaults whether `rand_impl` may be specialized
+# for the measure (coming from the default `rand_impl` itself, it is not):
+struct _NoRandImpl end
+struct _MaybeRandImpl end
 
-@inline function _rand_via_std(ctx::GenContext, μ, ::Type{S}, ::IntegerLike, ::Tuple{}) where {S<:StdMeasure}
-    convert_realtype(get_precision(ctx), transport_from_std(S, μ, rand_impl(ctx, S())))
-end
-@inline function _rand_via_std(ctx::GenContext, μ, ::Type{S}, n::IntegerLike, ::Any) where {S<:StdMeasure}
-    convert_realtype(get_precision(ctx), transport_from_std(S, μ, _rand_std(ctx, S, (dynamic(n),))))
-end
-@inline function _rand_via_std(ctx::GenContext, μ, ::Type{AnyStdMeasure}, n::IntegerLike, sz)
-    _rand_via_std(ctx, μ, StdUniform, n, sz)
-end
-function _rand_via_std(::GenContext, μ, ::Any, ::Any, ::Any)
-    throw(ArgumentError("Random variate generation is not implemented for measures of type $(nameof(typeof(μ))), define MeasureBase.rand_impl"))
-end
+@inline rand_impl(ctx::GenContext, μ) = _rand_default(ctx, μ, (), _NoRandImpl())
+@inline batched_rand_impl(ctx::GenContext, μ, sz::Dims) = _rand_default(ctx, μ, sz, _MaybeRandImpl())
 
+@inline _rand_default(ctx::GenContext, μ, sz::Dims, m) = _rand_via_std(ctx, μ, sz, preferred_stdmeasure(μ), m)
 
-"""
-    MeasureBase.batched_rand_impl(ctx::GenContext, μ, sz::Dims)
-
-Generate a batch of random variates of `μ` of batch size `sz` in flat
-form, an array of size `(flat variate dims..., sz...)` (see
-[`MeasureBase.mspace_flatsize`](@ref)). Measures with variates of
-unknown flat size only support batches of scalar variates.
-
-The default implementation draws a batch of variates of the preferred
-standard measure of `μ` and transports it to `μ`, or generates the
-variates one by one if `μ` has no standard transport.
-"""
-function batched_rand_impl end
-
-function batched_rand_impl(ctx::GenContext, μ, sz::Dims)
-    _batched_rand_via_std(ctx, μ, sz, preferred_stdmeasure(μ), fast_dof(μ), mspace_flatsize(μ))
+@inline function _rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{S}, m) where {S<:StdMeasure}
+    _rand_via_std_dof(ctx, μ, sz, S, fast_dof(μ), m)
 end
+@inline _rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{AnyStdMeasure}, m) = _rand_via_std(ctx, μ, sz, StdUniform, m)
+@inline _rand_via_std(ctx::GenContext, μ, sz::Dims, ::Any, m) = _rand_pointwise(ctx, μ, sz, m)
 
-function _batched_rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{S}, n::IntegerLike, ::SizeLike) where {S<:StdMeasure}
+function _rand_via_std_dof(ctx::GenContext, μ, sz::Dims, ::Type{S}, n::IntegerLike, ::Any) where {S<:StdMeasure}
     convert_realtype(get_precision(ctx), batched_transport_from_std(S, μ, _rand_std(ctx, S, (dynamic(n), sz...))))
 end
-@inline function _batched_rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{AnyStdMeasure}, n::IntegerLike, sz_flat::SizeLike)
-    _batched_rand_via_std(ctx, μ, sz, StdUniform, n, sz_flat)
-end
-@inline function _batched_rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{AnyStdMeasure}, n::IntegerLike, sz_flat::NoMSpaceElementSize)
-    _batched_rand_via_std(ctx, μ, sz, StdUniform, n, sz_flat)
-end
-function _batched_rand_via_std(ctx::GenContext, μ, sz::Dims, ::Any, ::Any, ::SizeLike)
-    _batched_rand_pointwise(ctx, μ, sz)
-end
-function _batched_rand_via_std(::GenContext, μ, ::Dims, ::Any, ::Any, ::NoMSpaceElementSize)
-    throw(ArgumentError("Batched random variate generation requires measures of type $(nameof(typeof(μ))) to have a known variate size"))
-end
-function _batched_rand_via_std(::GenContext, μ, ::Dims, ::Type{S}, ::IntegerLike, ::NoMSpaceElementSize) where {S<:StdMeasure}
-    throw(ArgumentError("Batched random variate generation requires measures of type $(nameof(typeof(μ))) to have a known variate size"))
+@inline _rand_via_std_dof(ctx::GenContext, μ, sz::Dims, ::Type, ::Any, m) = _rand_pointwise(ctx, μ, sz, m)
+
+# Variates generated one by one, stacked into a flat batch:
+@inline _rand_pointwise(ctx::GenContext, μ, sz::Dims, ::Any) = _batched_rand_pointwise(ctx, μ, sz)
+@inline _rand_pointwise(ctx::GenContext, μ, ::Tuple{}, ::_MaybeRandImpl) = rand_impl(ctx, μ)
+@noinline function _rand_pointwise(::GenContext, μ, ::Tuple{}, ::_NoRandImpl)
+    throw(ArgumentError("Random variate generation is not implemented for measures of type $(nameof(typeof(μ))), define MeasureBase.batched_rand_impl or MeasureBase.rand_impl"))
 end
 
 function _batched_rand_pointwise(ctx::GenContext, μ, sz::Dims)
     _stack_variates(map(_ -> rand_impl(ctx, μ), CartesianIndices(sz)))
 end
+@inline _batched_rand_pointwise(ctx::GenContext, μ, ::Tuple{}) = rand_impl(ctx, μ)
 
 @inline _stack_variates(xs::AbstractArray{<:Number}) = xs
-@inline _stack_variates(xs::AbstractArray{<:AbstractArray}) = stacked(xs)
+@inline _stack_variates(xs::AbstractArray{<:AbstractArray}) = stacked(map(_stack_variates, xs))
+@inline _stack_variates(xs::AbstractArray{<:Union{Tuple,NamedTuple}}) = StructArrays.components(StructArray(xs))
 
 
-# Bulk draws of standard variates on the compute unit:
+# Bulk draws of standard variates on the compute unit, single draws for
+# zero batch dimensions:
 
 @inline _rand_std(ctx::GenContext, ::Type{S}, dims::Dims) where {S<:StdMeasure} = batched_rand_impl(ctx, S(), dims)
 
@@ -115,17 +107,26 @@ end
 # Not all compute units provide exponential draws, derive them from uniform draws then:
 @inline _randexp_bulk(ctx::GenContext, sz::Dims, ::AbstractComputeUnit) = -log1p.(-_rand_bulk(ctx, sz))
 
+@inline _rand_bulk(ctx::GenContext, ::Tuple{}) = rand(get_rng(ctx), get_precision(ctx))
+@inline _randn_bulk(ctx::GenContext, ::Tuple{}) = randn(get_rng(ctx), get_precision(ctx))
+@inline _randexp_bulk(ctx::GenContext, ::Tuple{}) = randexp(get_rng(ctx), get_precision(ctx))
+
 # Test values use a constant RNG, which only draws single values:
 const _ConstantContext = GenContext{<:AbstractFloat,<:AbstractComputeUnit,ConstantRNG}
 @inline _rand_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, rand(ConstantRNG(), get_precision(ctx)), sz)
 @inline _randn_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, randn(ConstantRNG(), get_precision(ctx)), sz)
 @inline _randexp_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, randexp(ConstantRNG(), get_precision(ctx)), sz)
+@inline _rand_bulk(ctx::_ConstantContext, ::Tuple{}) = rand(ConstantRNG(), get_precision(ctx))
+@inline _randn_bulk(ctx::_ConstantContext, ::Tuple{}) = randn(ConstantRNG(), get_precision(ctx))
+@inline _randexp_bulk(ctx::_ConstantContext, ::Tuple{}) = randexp(ConstantRNG(), get_precision(ctx))
 @inline _const_bulk(ctx::GenContext, x, sz::Dims) = fill!(allocate_array(ctx, typeof(x), sz), x)
 
-# A mask over the batch dimensions, aligned with a flat batch of variates:
-@inline _batch_mask(mask::AbstractArray, ::Tuple{}) = mask
-@inline function _batch_mask(mask::AbstractArray, sz_flat::SizeLike)
-    reshape(mask, (ntuple(_ -> 1, Val(length(sz_flat)))..., size(mask)...))
+# A mask over the batch dimensions, aligned with a flat batch of variates
+# of rank `k`:
+@inline _batch_mask(mask::Number, ::Any) = mask
+@inline _batch_mask(mask::AbstractArray, ::StaticInteger{0}) = mask
+@inline function _batch_mask(mask::AbstractArray, ::StaticInteger{K}) where {K}
+    reshape(mask, (ntuple(_ -> 1, Val(K))..., size(mask)...))
 end
 
 # A batch of copies of a constant variate:
@@ -134,6 +135,5 @@ function _const_batch(ctx::GenContext, x, sz::Dims)
     X .= x
     return X
 end
-function _const_batch(ctx::GenContext, x::Number, sz::Dims)
-    fill!(allocate_array(ctx, typeof(x), sz), x)
-end
+@inline _const_batch(ctx::GenContext, x::Number, sz::Dims) = fill!(allocate_array(ctx, typeof(x), sz), x)
+@inline _const_batch(::GenContext, x::Number, ::Tuple{}) = x
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index ea767146..3c2ab19b 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -104,3 +104,4 @@ end
 
 @inline batched_transport_to_std(::Type{S}, ::S, X::AbstractArray) where {S<:StdMeasure} = _as_stdstream_batch(X)
 @inline batched_transport_from_std(::Type{S}, ::S, Z::AbstractArray) where {S<:StdMeasure} = _drop_stdstream_dim(Z)
+@inline batched_transport_from_std(::Type{S}, ::S, z::AbstractVector) where {S<:StdMeasure} = z[begin]
diff --git a/test/cuda/Project.toml b/test/cuda/Project.toml
new file mode 100644
index 00000000..b9159db5
--- /dev/null
+++ b/test/cuda/Project.toml
@@ -0,0 +1,13 @@
+[deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+AffineMaps = "2c83c9a8-abf5-4329-a0d7-deffaf474661"
+ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+HeterogeneousComputing = "2182be2a-124f-4a91-8389-f06db5907a21"
+MeasureBase = "fa1605e6-acd5-459c-a1e6-7e635759db14"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+MeasureBase = {path = "../.."}
diff --git a/test/cuda/runtests.jl b/test/cuda/runtests.jl
new file mode 100644
index 00000000..dd34cd79
--- /dev/null
+++ b/test/cuda/runtests.jl
@@ -0,0 +1,116 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+# CUDA tests, not part of the default test suite (they need a CUDA GPU).
+# Run with `julia --project=test/cuda test/cuda/runtests.jl` after
+# instantiating that project.
+
+using Test
+using CUDA
+using Adapt: adapt
+using HeterogeneousComputing: GenContext, AbstractComputeUnit
+using MeasureBase
+using MeasureBase: StdNormal, StdUniform, StdExponential, Dirac
+using MeasureBase: productmeasure, pushfwd, mcombine, weightedmeasure, superpose, SpikeMixture
+using MeasureBase: batched_rand_impl
+using MeasureBase.InverseFunctions: inverse
+using ArraysOfArrays: sliced, flatview
+using AffineMaps: Mul, MulAdd
+using Distributions: Normal
+
+CUDA.allowscalar(false)
+
+# Evaluates `f` on device copies of `args` and compares with the plain
+# result, the result must live on the device:
+function test_cuda(f, args...)
+    expected = f(args...)
+    result = f(map(cu_copy, args)...)
+    @test _device_array(result)
+    @test _plain(result) ≈ _plain(expected) nans = true
+    return result
+end
+
+cu_copy(x::AbstractArray) = CuArray(x)
+cu_copy(μ::AbstractMeasure) = adapt(CuArray, μ)
+cu_copy(x) = x
+_device_array(x::AbstractArray) = parent_array(x) isa CuArray
+_device_array(x::Tuple) = all(_device_array, x)
+parent_array(x::CuArray) = x
+parent_array(x::AbstractArray) = parent_array(parent(x))
+parent_array(x::Base.ReshapedArray) = parent_array(parent(x))
+_plain(x::AbstractArray) = Array(flatview(x))
+_plain(x::Tuple) = map(_plain, x)
+
+@testset "CUDA" begin
+    X = randn(3, 20)
+    Xc = vcat(randn(2, 20), rand(1, 20))
+
+    @testset "densities" begin
+        test_cuda(X -> logdensities(StdNormal()^3, X), X)
+        test_cuda(X -> logdensities(StdNormal()^3, sliced(X, Val(1))), X)
+        test_cuda(X -> logdensities((StdNormal()^3)^4, reshape(X[:, 1:16], 3, 4, 4)), X)
+        test_cuda(X -> logdensities(weightedmeasure(log(0.3), StdNormal()^3), X), X)
+        mc = mcombine(vcat, StdNormal()^2, StdUniform()^1)
+        test_cuda(X -> logdensities(mc, X), Xc)
+        mix = superpose(weightedmeasure(log(0.3), StdNormal()), weightedmeasure(log(0.7), StdUniform()))
+        test_cuda(x -> logdensities(mix, x), rand(20))
+        test_cuda(x -> logdensities(SpikeMixture(StdNormal(), 0.2), x), vcat(randn(19), 0.0))
+        test_cuda(x -> logdensities(Dirac(0.5), x), vcat(rand(19), 0.5))
+        νe = pushfwd(Base.BroadcastFunction(exp), StdNormal()^3)
+        test_cuda(Y -> logdensities(νe, Y), exp.(X))
+        P = productmeasure([pushfwd(Mul(s), StdNormal()) for s in (1.0, 2.0, 3.0)])
+        test_cuda((P, X) -> logdensities(P, X), P, X)
+        test_cuda((P, X) -> logdensities(P, sliced(X, Val(1))), P, X)
+        Pa = productmeasure([MeasureBase.AsMeasure{Normal{Float64}}(Normal(μ, 1.0)) for μ in (0.0, 1.0, 2.0)])
+        test_cuda((P, X) -> logdensities(P, X), Pa, X)
+        # Array products of array-variate marginals loop over the marginals
+        # on the host, the marginals stay host arrays:
+        Pv = productmeasure([weightedmeasure(log(i), StdNormal()^2) for i in 1:3])
+        test_cuda(X -> logdensities(Pv, X), randn(2, 3, 5))
+    end
+
+    @testset "transports" begin
+        g = transport_to(StdUniform()^3, StdNormal()^3)
+        test_cuda(X -> flatview(g.(sliced(X, Val(1)))), X)
+        test_cuda(X -> flatview(inverse(g).(g.(sliced(X, Val(1))))), X)
+        h = transport_to(StdUniform()^(2, 3), (StdNormal()^2)^3)
+        test_cuda(X -> flatview(h.(sliced(X, Val(2)))), randn(2, 3, 4))
+        mc = mcombine(vcat, StdNormal()^2, StdUniform()^1)
+        c = transport_to(StdExponential()^3, mc)
+        test_cuda(X -> flatview(c.(sliced(X, Val(1)))), Xc)
+        test_cuda(X -> flatview(inverse(c).(c.(sliced(X, Val(1))))), Xc)
+        P = productmeasure([pushfwd(Mul(s), StdNormal()) for s in (1.0, 2.0, 3.0)])
+        f = transport_to(StdUniform()^3, P)
+        test_cuda((P, X) -> flatview(transport_to(StdUniform()^3, P).(sliced(X, Val(1)))), P, X)
+        test_cuda((P, X) -> flatview(transport_to(P, StdUniform()^3).(sliced(X, Val(1)))), P, rand(3, 20))
+        Pv = productmeasure([weightedmeasure(log(i), StdNormal()^2) for i in 1:3])
+        fv = transport_to(StdUniform()^6, Pv)
+        test_cuda(X -> flatview(fv.(sliced(X, Val(2)))), randn(2, 3, 5))
+        test_cuda(X -> flatview(inverse(fv).(fv.(sliced(X, Val(2))))), randn(2, 3, 5))
+        νe = pushfwd(Base.BroadcastFunction(exp), StdNormal()^3)
+        fe = transport_to(StdUniform()^3, νe)
+        test_cuda(Y -> flatview(fe.(sliced(Y, Val(1)))), exp.(X))
+        A = [2.0 0.5; 0.0 1.5]
+        b = [1.0, -1.0]
+        νa = pushfwd(MulAdd(A, b), StdNormal()^2)
+        νac = pushfwd(MulAdd(CuArray(A), CuArray(b)), StdNormal()^2)
+        Ya = randn(2, 20)
+        ra = flatview(transport_to(StdUniform()^2, νac).(sliced(CuArray(Ya), Val(1))))
+        @test ra isa CuArray && Array(ra) ≈ flatview(transport_to(StdUniform()^2, νa).(sliced(Ya, Val(1))))
+        # AffineMaps has no device support for log-abs-det-Jacobians yet:
+        @test_broken Array(logdensities(νac, CuArray(Ya))) ≈ logdensities(νa, Ya)
+    end
+
+    @testset "random variates" begin
+        ctx = GenContext{Float32}(AbstractComputeUnit(CUDA.device()), CUDA.default_rng())
+        for μ in (StdNormal(), StdUniform(), StdExponential(), StdNormal()^3, (StdNormal()^2)^3, weightedmeasure(0.3, StdNormal()^2), mcombine(vcat, StdNormal()^2, StdUniform()^1), pushfwd(Base.BroadcastFunction(exp), StdNormal()^2), superpose(StdNormal(), StdUniform()), SpikeMixture(StdNormal(), 0.5))
+            X = batched_rand_impl(ctx, μ, (100,))
+            @test X isa CuArray{Float32}
+            ℓ = logdensities(μ, X)
+            @test ℓ isa CuArray && all(isfinite, Array(ℓ))
+        end
+        Pt = productmeasure((StdNormal(), StdExponential()^2))
+        Xt = batched_rand_impl(ctx, Pt, (50,))
+        @test Xt isa Tuple && all(x -> x isa CuArray{Float32}, Xt)
+        @test size(Xt[2]) == (2, 50)
+    end
+end
diff --git a/test/rand_batched.jl b/test/rand_batched.jl
new file mode 100644
index 00000000..ee63ef75
--- /dev/null
+++ b/test/rand_batched.jl
@@ -0,0 +1,102 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using Random, Statistics
+using StableRNGs: StableRNG
+using StructArrays: StructArray
+using ArraysOfArrays: flatview
+
+using MeasureBase
+using MeasureBase: GenContext
+using MeasureBase: StdUniform, StdExponential, StdLogistic, StdNormal, Dirac
+using MeasureBase: weightedmeasure, superpose, mcombine, mbind, productmeasure, pushfwd, SpikeMixture
+using MeasureBase: rand_impl, batched_rand_impl
+using Distributions: Normal, MvNormal, logpdf
+
+struct NoBatchRandMeasure <: AbstractMeasure end
+MeasureBase.rand_impl(ctx::GenContext, ::NoBatchRandMeasure) = 3 * rand(MeasureBase.get_rng(ctx), MeasureBase.get_precision(ctx))
+
+struct NoRandMeasure <: AbstractMeasure end
+
+@testset "batched rand" begin
+    stblrng() = StableRNG(789)
+    ctx() = GenContext{Float64}(stblrng())
+
+    @testset "single variates as batches with zero batch dimensions" begin
+        @test @inferred(batched_rand_impl(ctx(), StdNormal(), ())) isa Float64
+        @test batched_rand_impl(ctx(), StdNormal(), ()) == rand(stblrng(), StdNormal())
+        @test @inferred(batched_rand_impl(ctx(), StdExponential(), ())) isa Float64
+        @test @inferred(batched_rand_impl(ctx(), StdLogistic(), ())) isa Float64
+        @test @inferred(batched_rand_impl(ctx(), StdUniform(), ())) isa Float64
+        @test @inferred(batched_rand_impl(ctx(), StdNormal()^3, ())) isa Vector{Float64}
+        @test batched_rand_impl(ctx(), StdNormal()^3, ()) == rand(stblrng(), StdNormal()^3)
+        @test batched_rand_impl(ctx(), (StdNormal()^2)^3, ()) == flatview(rand(stblrng(), (StdNormal()^2)^3))
+        @test batched_rand_impl(ctx(), Dirac(2.5), ()) === 2.5
+        @test batched_rand_impl(ctx(), Dirac([1.0, 2.0]), ()) == [1.0, 2.0]
+        @test batched_rand_impl(ctx(), weightedmeasure(0.3, StdNormal()), ()) isa Float64
+        @test batched_rand_impl(ctx(), pushfwd(exp, StdNormal()), ()) isa Float64
+        @test batched_rand_impl(ctx(), pushfwd(Base.BroadcastFunction(exp), StdNormal()^2), ()) isa Vector{Float64}
+        @test batched_rand_impl(ctx(), SpikeMixture(StdNormal(), 0.5), ()) isa Float64
+        @test batched_rand_impl(ctx(), superpose(StdNormal(), StdUniform()), ()) isa Float64
+        @test batched_rand_impl(ctx(), NoBatchRandMeasure(), ()) isa Float64
+        @test batched_rand_impl(ctx(), NoBatchRandMeasure(), (4,)) isa Vector{Float64}
+        @test_throws ArgumentError rand(NoRandMeasure())
+        @test_throws ArgumentError batched_rand_impl(ctx(), NoRandMeasure(), ())
+        @test_throws ArgumentError batched_rand_impl(ctx(), NoRandMeasure(), (3,))
+    end
+
+    @testset "structured batches" begin
+        Pt = productmeasure((StdNormal(), StdExponential()^2))
+        Xt = batched_rand_impl(ctx(), Pt, (5,))
+        @test Xt isa Tuple && size(Xt[1]) == (5,) && size(Xt[2]) == (2, 5)
+        xt = batched_rand_impl(ctx(), Pt, ())
+        @test xt isa Tuple{Float64,Vector{Float64}} && xt == rand(stblrng(), Pt)
+        XPt = rand(stblrng(), Pt^5)
+        @test XPt isa StructArray && length(XPt) == 5
+        @test XPt[2] isa Tuple{Float64,<:AbstractVector{Float64}} && length(XPt[2][2]) == 2
+        @test logdensityof(Pt^5, collect(XPt)) ≈ sum(logdensityof.(Ref(Pt), XPt))
+        Pn = productmeasure((a = StdNormal(), b = StdExponential()^2))
+        Xn = batched_rand_impl(ctx(), Pn, (2, 3))
+        @test Xn isa NamedTuple{(:a, :b)} && size(Xn.a) == (2, 3) && size(Xn.b) == (2, 2, 3)
+        XPn = rand(stblrng(), Pn^4)
+        @test XPn isa StructArray && XPn[1] isa NamedTuple{(:a, :b)}
+
+        mm = mcombine(merge, productmeasure((a = StdNormal(),)), productmeasure((b = StdUniform()^2,)))
+        Xm = batched_rand_impl(ctx(), mm, (3,))
+        @test Xm isa NamedTuple{(:a, :b)} && size(Xm.a) == (3,) && size(Xm.b) == (2, 3)
+        @test batched_rand_impl(ctx(), mm, ()) == rand(stblrng(), mm)
+        mt = mcombine(tuple, StdNormal(), StdUniform()^2)
+        Xtt = batched_rand_impl(ctx(), mt, (3,))
+        @test Xtt isa Tuple && size(Xtt[1]) == (3,) && size(Xtt[2]) == (2, 3)
+        @test batched_rand_impl(ctx(), mt, ()) == rand(stblrng(), mt)
+    end
+
+    @testset "powers and value-dependent sizes" begin
+        f_β(a) = StdNormal()^length(a)
+        μb = mbind(f_β, StdUniform()^1, vcat)
+        x = rand(stblrng(), μb^3)
+        @test x isa Vector{Vector{Float64}} && length(x) == 3 && all(length.(x) .== 2)
+        X = batched_rand_impl(ctx(), μb^3, (4,))
+        @test size(X) == (2, 3, 4)
+        @test batched_rand_impl(ctx(), μb, ()) == rand(stblrng(), μb)
+        @test size(batched_rand_impl(ctx(), (StdNormal()^2)^3, (4, 5))) == (2, 3, 4, 5)
+        @test size(batched_rand_impl(ctx(), StdNormal()^(2, 3), (4,))) == (2, 3, 4)
+    end
+
+    @testset "moments of batches" begin
+        n = 20_000
+        X = batched_rand_impl(ctx(), superpose(StdNormal(), Dirac(3.0)), (n,))
+        @test isapprox(mean(X), 1.5, atol = 0.05)
+        Xs = batched_rand_impl(ctx(), SpikeMixture(StdNormal()^2, 0.5), (n,))
+        @test size(Xs) == (2, n) && isapprox(mean(Xs .== 0), 0.5, atol = 0.02)
+        Xp = batched_rand_impl(ctx(), pushfwd(Base.BroadcastFunction(exp), StdNormal()^2), (n,))
+        @test isapprox(mean(log.(Xp)), 0.0, atol = 0.03)
+        Xd = batched_rand_impl(ctx(), MeasureBase.AsMeasure{Normal{Float64}}(Normal(1.0, 2.0)), (n,))
+        @test Xd isa Vector{Float64} && isapprox(mean(Xd), 1.0, atol = 0.05)
+        xd = batched_rand_impl(ctx(), MeasureBase.AsMeasure{Normal{Float64}}(Normal(1.0, 2.0)), ())
+        @test xd isa Float64
+        Xmv = batched_rand_impl(ctx(), MeasureBase.AsMeasure{typeof(MvNormal([1.0, 2.0], [1.0, 0.5]))}(MvNormal([1.0, 2.0], [1.0, 0.5])), (n,))
+        @test size(Xmv) == (2, n) && isapprox(vec(mean(Xmv, dims = 2)), [1.0, 2.0], atol = 0.05)
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 5a427976..d9071473 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -38,6 +38,7 @@ include("combinators/bind.jl")
 include("combinators/product.jl")
 
 include("rand.jl")
+include("rand_batched.jl")
 
 include("distributions/test_distributions.jl")
 

From f30ab14d35f8ac5daa3c38b0684869812137ec29 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 12:26:00 +0200
Subject: [PATCH 096/122] Evaluate structured batches and drop the last
 size-based routing

Batches of tuple and named tuple variates enter as struct arrays or
arrays of tuples: their flat storage is the tuple of the flat component
storages, which the batched kernels of products consume directly, and
transport functions broadcast over such batches and return struct
arrays for tuple-valued targets. Support checks of powers route by the
variate rank of the innermost base measure. The Reactant runner takes
its backend from MEASUREBASE_REACTANT_BACKEND and materializes array
results inside the compiled functions, so it also runs on GPUs.

Created by generative AI.
---
 src/combinators/power.jl   | 15 +++++--------
 src/combinators/product.jl |  4 +++-
 src/density-batched.jl     | 19 +++++++++++++++-
 src/transport-batched.jl   | 15 ++++++++++++-
 test/reactant/runtests.jl  | 14 ++++++++----
 test/runtests.jl           |  1 +
 test/structured_batches.jl | 45 ++++++++++++++++++++++++++++++++++++++
 7 files changed, 97 insertions(+), 16 deletions(-)
 create mode 100644 test/structured_batches.jl

diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index baaf033e..e29aa98a 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -152,7 +152,7 @@ end
     _check_pwr_shape(μ, x)
     _powered_point_nested(f, μ, x, _flat_storage(x))
 end
-@inline function _powered_point_nested(f::F, μ::PowerMeasure, x, x_flat::AbstractArray) where {F}
+@inline function _powered_point_nested(f::F, μ::PowerMeasure, x, x_flat::Union{AbstractArray,Tuple,NamedTuple}) where {F}
     _point_result(_materialize(_batched_kernel(f, μ, x_flat)), μ)
 end
 function _powered_point_nested(f::F, μ::PowerMeasure, x::AbstractArray, ::NoFlatStorage) where {F}
@@ -182,20 +182,17 @@ function _powered_ld_with_rest(μ::PowerMeasure, X::AbstractArray, sz::Dims)
     return _sum_leading_dims(ℓ, static(length(pwr_axes(μ)))), X_rest
 end
 
-# Support checks of powers run over the flat variate storage where the base
-# measure has scalar variates, elementwise otherwise:
+# Support checks of powers run over the flat variate storage where the
+# innermost base measure has scalar variates, elementwise otherwise:
 @inline function insupport(μ::PowerMeasure, x::AbstractArray)
-    _powered_insupport(μ, x, _flat_storage(x), mspace_flatsize(μ))
+    ν, _ = _pwr_unwrap(μ)
+    _powered_insupport(μ, x, _flat_storage(x), _static_ndims(ν))
 end
 
-@inline function _powered_insupport(μ::PowerMeasure, x, x_flat::AbstractArray, ::SizeLike)
+@inline function _powered_insupport(μ::PowerMeasure, x, x_flat::AbstractArray, ::StaticInteger{0})
     ν, _ = _pwr_unwrap(μ)
-    _powered_insupport_flat(ν, x_flat, mspace_flatsize(ν))
-end
-@inline function _powered_insupport_flat(ν, x_flat::AbstractArray, ::Tuple{})
     _all_insupport(broadcast(_insupport_bool ∘ Base.Fix1(insupport, ν), x_flat))
 end
-@inline _powered_insupport_flat(ν, x_flat::AbstractArray, ::Any) = _powered_insupport_elementwise(ν, x_flat)
 @inline _powered_insupport(μ::PowerMeasure, x, ::Any, ::Any) = _powered_insupport_elementwise(pwr_base(μ), x)
 
 @inline function _powered_insupport_elementwise(ν, x::AbstractArray)
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index bed70a5b..b134e067 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -28,7 +28,7 @@ Base.size(μ::AbstractProductMeasure) = size(marginals(μ))
 
 basemeasure(d::AbstractProductMeasure) = productmeasure(map(basemeasure, marginals(d)))
 
-rand_impl(ctx::GenContext, d::AbstractProductMeasure) = map(Base.Fix1(_marginal_rand, ctx), marginals(d))
+rand_impl(ctx::GenContext, d::AbstractProductMeasure) = _map(Base.Fix1(_marginal_rand, ctx), marginals(d))
 
 @inline _marginal_rand(ctx::GenContext, m::AbstractMeasure) = rand_impl(ctx, m)
 @inline _marginal_rand(ctx::GenContext, d) = convert_realtype(get_precision(ctx), rand(get_rng(ctx), d))
@@ -281,6 +281,8 @@ end
 # TODO: Better `map` support in MappedArrays
 _map(f, args...) = map(f, args...)
 _map(f, x::MappedArrays.ReadonlyMappedArray) = mappedarray(fchain((x.f, f)), x.data)
+# Variates of struct array marginals are collected into plain arrays:
+_map(f, x::StructArray) = map(f, collect(x))
 
 function testvalue(::Type{T}, d::AbstractProductMeasure) where {T}
     _map(m -> testvalue(T, m), marginals(d))
diff --git a/src/density-batched.jl b/src/density-batched.jl
index bb28947f..ea6e1f99 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -114,7 +114,9 @@ end
 @inline _batched_kernel(::typeof(logdensity_def), μ, X) = batched_logdensity_def(μ, X)
 
 # Flat storage of a (nested) batch: the underlying array of memory-ordered
-# split arrays, a stacked copy for other known split modes.
+# split arrays, a stacked copy for other known split modes. Struct arrays
+# of tuple and named tuple variates have the flat storage of their
+# components (copied where necessary, the batch dimensions are shared).
 struct NoFlatStorage end
 @inline _flat_storage(X::AbstractArray{<:Number}) = X
 @inline _flat_storage(X::AbstractArray) = _flat_storage_bymode(X, getsplitmode(X))
@@ -125,6 +127,18 @@ end
 @inline _flat_storage_bymode(::AbstractArray, ::UnknownSplitMode) = NoFlatStorage()
 @inline _flat_storage_bymode(::AbstractArray, ::NonSplitMode) = NoFlatStorage()
 
+@inline _flat_storage(X::StructArray{<:Union{Tuple,NamedTuple}}) = _components_storage(StructArrays.components(X))
+@inline function _flat_storage(X::AbstractArray{<:Union{Tuple,NamedTuple}})
+    _flat_storage(StructArray(X; unwrap = T -> T <: Union{Tuple,NamedTuple}))
+end
+@inline _components_storage(cs::Tuple) = map(_component_storage, cs)
+@inline _components_storage(cs::NamedTuple{names}) where {names} = NamedTuple{names}(map(_component_storage, values(cs)))
+@inline _component_storage(c::StructArray{<:Union{Tuple,NamedTuple}}) = _flat_storage(c)
+@inline _component_storage(c::AbstractArray{<:Number}) = c
+@inline _component_storage(c::AbstractArray{<:AbstractArray}) = _component_flat(c, _flat_storage(c))
+@inline _component_flat(c, c_flat::AbstractArray) = c_flat
+@inline _component_flat(c, ::NoFlatStorage) = stacked(c)
+
 # Entry: arrays of numbers are flat storage, arrays of variates are fused
 # into their flat storage (else evaluated variate by variate), tuples and
 # named tuples of batches go to the kernels directly.
@@ -134,6 +148,9 @@ end
 @inline function _batched_ld_nested(f::F, μ, X::AbstractArray, X_flat::AbstractArray, ::StaticInteger) where {F}
     _check_batch_shape(_batched_kernel(f, μ, X_flat), X)
 end
+@inline function _batched_ld_nested(f::F, μ, X::AbstractArray, X_flat::Union{Tuple,NamedTuple}, ::Any) where {F}
+    _check_batch_shape(_batched_kernel(f, μ, X_flat), X)
+end
 @inline function _batched_ld_nested(f::F, μ, X::AbstractArray, ::Any, ::Any) where {F}
     Broadcast.instantiate(Broadcast.broadcasted(_PointLogd(f, μ), X))
 end
diff --git a/src/transport-batched.jl b/src/transport-batched.jl
index b901478a..40bd9c56 100644
--- a/src/transport-batched.jl
+++ b/src/transport-batched.jl
@@ -243,6 +243,17 @@ function _broadcast_transport(f::TransportFunction, X, X_flat::AbstractArray, ::
     return _batch_variates(Y_flat, f.ν, Val(K))
 end
 
+# Batches of tuple and named tuple variates are tuples of batches, the
+# target layout follows from the target measure:
+function _broadcast_transport(f::TransportFunction, X, X_flat::Union{Tuple,NamedTuple}, ::Any, ::Any)
+    _structured_variates(batched_transport_def(f.ν, f.μ, X_flat), f.ν)
+end
+function _broadcast_transport(f::TransportFunction{<:ProductMeasure{<:Union{Tuple,NamedTuple}}}, X, X_flat::AbstractArray, ::StaticInteger, ::NoMSpaceElementSize)
+    _structured_variates(batched_transport_def(f.ν, f.μ, X_flat), f.ν)
+end
+@inline _structured_variates(Y::Union{Tuple,NamedTuple}, ν) = _pwr_variate(ν, Y)
+@inline _structured_variates(Y::AbstractArray, ν) = _batch_variates(Y, ν, Val(dynamic(_static_ndims(ν))))
+
 _broadcast_transport(f::TransportFunction, X, ::Any, ::Any, ::Any) = map(_Pointwise(f), X)
 
 # Prevents re-entering the broadcast hook from `map` implementations that
@@ -253,8 +264,10 @@ end
 @inline (p::_Pointwise)(x) = p.f(x)
 
 # The batch of variates in the layout of the target measure over the flat
-# result, nested powers included:
+# result, nested powers included, batches of tuple variates as struct
+# arrays:
 @inline _batch_variates(Y::AbstractArray, ν, ::Val{K}) where {K} = _nest_batch(Y, Val(K))
+@inline _batch_variates(Y::Union{Tuple,NamedTuple}, ν, ::Val) = _pwr_variate(ν, Y)
 @inline function _batch_variates(Y::AbstractArray, ν::PowerMeasure, ::Val)
     sliced(_pwr_variate(ν, Y), Val(length(pwr_axes(ν))))
 end
diff --git a/test/reactant/runtests.jl b/test/reactant/runtests.jl
index 1357cf8d..0fee2da1 100644
--- a/test/reactant/runtests.jl
+++ b/test/reactant/runtests.jl
@@ -3,7 +3,8 @@
 # Reactant smoke tests, not part of the default test suite. Run with
 # `julia --project=test/reactant test/reactant/runtests.jl` after
 # instantiating that project, or include this file in an environment that
-# provides Reactant.
+# provides Reactant. The backend defaults to the CPU, set the environment
+# variable `MEASUREBASE_REACTANT_BACKEND` (e.g. to "gpu") to change it.
 
 using Test
 using Reactant
@@ -14,17 +15,22 @@ using MeasureBase: mcombine
 using ArraysOfArrays: VectorOfSimilarVectors, sliced, flatview
 using Distributions: Normal, Exponential, Uniform, Beta
 
-Reactant.set_default_backend("cpu")
+Reactant.set_default_backend(get(ENV, "MEASUREBASE_REACTANT_BACKEND", "cpu"))
 
-# Compiles `f` for traced copies of `args` and compares with the plain result:
+# Compiles `f` for traced copies of `args` and compares with the plain
+# result. Array results are copied inside the compiled function, so that
+# views and reshapes of device arrays come back as plain device arrays:
 function test_traced(f, args...; kwargs...)
     expected = f(args...)
     traced_args = map(Reactant.to_rarray, args)
-    result = @jit f(traced_args...)
+    g = (xs...) -> _contiguous(f(xs...))
+    result = @jit g(traced_args...)
     @test _plain(result) ≈ _plain(expected) nans = true
     return result
 end
 
+_contiguous(x::AbstractArray) = copy(x)
+_contiguous(x) = x
 _plain(x::AbstractArray) = Array(x)
 _plain(x::Number) = Float64(x)
 
diff --git a/test/runtests.jl b/test/runtests.jl
index d9071473..4945ad06 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -17,6 +17,7 @@ include("test_basics.jl")
 include("getdof.jl")
 include("shape_contract.jl")
 include("logdensities.jl")
+include("structured_batches.jl")
 include("numtype.jl")
 include("transport.jl")
 include("transport_batched.jl")
diff --git a/test/structured_batches.jl b/test/structured_batches.jl
new file mode 100644
index 00000000..184f01da
--- /dev/null
+++ b/test/structured_batches.jl
@@ -0,0 +1,45 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using MeasureBase
+using MeasureBase: StdNormal, StdUniform, StdExponential, productmeasure, insupport
+using MeasureBase.InverseFunctions: inverse
+using StructArrays: StructArray
+using ArraysOfArrays: flatview
+
+@testset "structured batches" begin
+    Pt = productmeasure((StdNormal(), StdExponential()^2))
+    xs = [rand(Pt) for _ in 1:6]
+    ℓ_ref = logdensityof.(Ref(Pt), xs)
+    @test logdensities(Pt, xs) ≈ ℓ_ref
+    @test logdensities(Pt, StructArray(xs)) ≈ ℓ_ref
+    X = rand(Pt^6)
+    @test X isa StructArray
+    @test logdensities(Pt, X) ≈ logdensityof.(Ref(Pt), X)
+    @test logdensityof(Pt^6, X) ≈ sum(logdensityof.(Ref(Pt), X))
+    @test logdensityof(Pt^6, collect(X)) ≈ logdensityof(Pt^6, X)
+    Xm = rand(Pt^(2, 3))
+    @test size(Xm) == (2, 3) && logdensityof(Pt^(2, 3), Xm) ≈ sum(logdensityof.(Ref(Pt), Xm))
+    @test_throws ArgumentError logdensityof(Pt^5, X)
+
+    Pn = productmeasure((a = StdNormal(), b = StdExponential()^2))
+    Xn = rand(Pn^5)
+    @test logdensities(Pn, Xn) ≈ logdensityof.(Ref(Pn), Xn)
+    @test logdensities(Pn, collect(Xn)) ≈ logdensityof.(Ref(Pn), Xn)
+
+    f = transport_to(StdUniform()^3, Pt)
+    Y = f.(X)
+    @test Y isa AbstractVector && length(Y) == 6 && flatview(Y) ≈ stack(map(f, X))
+    Xr = inverse(f).(Y)
+    @test Xr isa StructArray && all(map((a, b) -> all(map(≈, a, b)), Xr, X))
+    h = transport_to(Pn, Pt)
+    Yn = h.(X)
+    @test Yn isa StructArray && Yn[1] isa NamedTuple{(:a, :b)}
+    @test all(Yn[i].a ≈ X[i][1] && Yn[i].b ≈ X[i][2] for i in 1:6)
+    @test inverse(h).(Yn) isa StructArray
+
+    @test insupport(StdUniform()^3, [0.1, 0.5, 0.9]) && !insupport(StdUniform()^3, [0.1, 1.5, 0.9])
+    @test insupport((StdUniform()^2)^3, rand(2, 3)) && !insupport((StdUniform()^2)^3, fill(2.0, 2, 3))
+    @test insupport(StdUniform()^3, [0.1, 0.5, 0.9]) isa Bool
+end

From a2ff84c6b4b4396fcebf1bb0829c8f256f369fbd Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 13:01:48 +0200
Subject: [PATCH 097/122] Fix batched-first review findings

Function objects stay opaque columns of struct array marginals, since
their wrappers can't be rebuilt from fields. Support checks of powers
with array-variate bases map over the variate slices, Dirac kernels and
power collapses are restricted to numeric variates, and array variates
of measures without a declared rank go through the point kernel again.
Streams: powers of bases without fixed sizes consume single streams
element by element and refuse batches, tuple products get a stream
length, multiplicity forms and a marginal-wise point transport, powers of
tuple products treat numeric flat variates as streams and accept tuples
of batches, and vcat-combined variates are flattened. Static streams
keep static row counts, so combined measures of static powers are
allocation-free again. Ragged batches evaluate variate by variate,
weighted measures report missing ranks, some_dof no longer recurses
forever, binds have test values, and Distributions evaluate flat batches
of array variates directly.

Created by generative AI.
---
 Project.toml                                  |   6 +-
 .../distribution_measure.jl                   |   8 ++
 src/MeasureBase.jl                            |   2 +-
 src/combinators/bind.jl                       |  18 ++-
 src/combinators/combined.jl                   |  36 +++--
 src/combinators/power.jl                      |  92 +++++++++++-
 src/combinators/product.jl                    |  62 +++++++-
 src/combinators/smart-constructors.jl         |   9 +-
 src/combinators/weighted.jl                   |   2 +-
 src/density-batched.jl                        |  69 +++++----
 src/getdof.jl                                 |   8 +-
 src/primitives/dirac.jl                       |  16 ++-
 src/transport-batched.jl                      |  46 +++---
 test/batched_regressions.jl                   | 134 ++++++++++++++++++
 test/logdensities.jl                          |   2 +-
 test/runtests.jl                              |   1 +
 16 files changed, 429 insertions(+), 82 deletions(-)
 create mode 100644 test/batched_regressions.jl

diff --git a/Project.toml b/Project.toml
index 182c9cdc..5d59432f 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,8 +4,8 @@ version = "0.14.12"
 authors = ["Chad Scherrer <chad.scherrer@gmail.com>", "Oliver Schulz <oschulz@mpp.mpg.de>", "contributors"]
 
 [deps]
-ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 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"
@@ -32,9 +32,9 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
-StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 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"
 
@@ -66,7 +66,6 @@ MeasureBaseReactantExt = "Reactant"
 [compat]
 Adapt = "3.7, 4"
 AffineMaps = "0.3"
-StructArrays = "0.6.18, 0.7"
 ArgCheck = "1, 2"
 ArraysOfArrays = "1.3"
 ChainRulesCore = "1"
@@ -103,6 +102,7 @@ Static = "0.8, 1"
 StaticArrays = "1.5"
 StaticThings = "0.2"
 Statistics = "1"
+StructArrays = "0.6.18, 0.7"
 StatsBase = "0.33, 0.34"
 StatsFuns = "0.9, 1, 2"
 Test = "1"
diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index d2ecc0bf..b480e11e 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -47,6 +47,14 @@ end
 @inline DensityInterface.logdensityof(m::DistributionMeasure) = logdensityof(m.obj)
 
 @inline MeasureBase.logdensity_def(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x)
+
+# Distributions evaluate flat batches of array variates (the trailing
+# dimensions are batch dimensions) directly:
+for bhead in (:batched_logdensityof_impl, :batched_logdensity_def)
+    @eval function MeasureBase.$bhead(m::DistributionMeasure{<:ArrayLikeVariate{N}}, X::AbstractArray{<:Real}) where {N}
+        Distributions.logpdf(m.obj, X)
+    end
+end
 @inline MeasureBase.unsafe_logdensityof(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x)
 @inline MeasureBase.insupport(m::DistributionMeasure, x) = Distributions.insupport(m.obj, x)
 
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index de4d9528..a0d29857 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -65,7 +65,7 @@ using HeterogeneousComputing:
 using ArraysOfArrays:
     ArrayOfSimilarArrays, VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays,
     VectorOfSimilarVectors, flatview, fused, stacked, sliced, getsplitmode,
-    is_memordered_splitmode, AbstractSplitMode, UnknownSplitMode, NonSplitMode
+    is_memordered_splitmode, AbstractSplitMode, AbstractPartMode, UnknownSplitMode, NonSplitMode
 
 using OneTwoMany: firstarg, secondarg
 
diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl
index e6682b67..87c43bac 100644
--- a/src/combinators/bind.jl
+++ b/src/combinators/bind.jl
@@ -281,7 +281,11 @@ rootmeasure(::Bind) =
 
 basemeasure(::Bind) = throw(ArgumentError("basemeasure is not available for Bind"))
 
-testvalue(::Bind) = throw(ArgumentError("testvalue is not available for Bind"))
+# Test values follow the primary test value through the secondary measure:
+function testvalue(::Type{T}, μ::Bind) where {T}
+    a = testvalue(T, μ.α)
+    _combine_variates(μ.f_c, a, testvalue(T, _get_β_a(μ, a)))
+end
 
 logdensity_def(::Bind, x) =
     throw(ArgumentError("logdensity_def is not available for Bind"))
@@ -336,19 +340,25 @@ function logdensityof_with_rest(μ::_BindBy{typeof(merge)}, x::NamedTuple)
     return ℓ_a + ℓ_b, merge(a, b), x_rest
 end
 
-function batched_logdensityof_with_rest(μ::_BindBy{typeof(vcat)}, x::AbstractVector, ::Tuple{})
+function batched_logdensityof_with_rest(μ::Bind, x::AbstractVector, ::Tuple{})
     ℓ, _, x_rest = logdensityof_with_rest(μ, x)
     return ℓ, x_rest
 end
 
 batched_logdensityof_impl(μ::_BindBy{typeof(vcat)}, X::AbstractArray) = _streamwise_ld(logdensityof_impl, μ, X)
+
+# Batches of streams containing binds are consumed stream by stream (by
+# the outermost stream combinator, see `fixed_stream_size`):
+@noinline function batched_logdensityof_with_rest(::Bind, ::AbstractArray, ::Dims)
+    throw(ArgumentError("Batches of variate streams containing binds must be consumed stream by stream"))
+end
 batched_logdensityof_impl(μ::_BindBy{typeof(vcat)}, x::AbstractVector) = _bind_ld_impl(vcat, μ, x)
 
 
 function rand_impl(ctx::GenContext, μ::Bind)
     a = rand_impl(ctx, μ.α)
     b = rand_impl(ctx, _get_β_a(μ, a))
-    return μ.f_c(a, b)
+    return _combine_variates(μ.f_c, a, b)
 end
 
 # The secondary measure depends on the primary variate, so batches are
@@ -390,7 +400,7 @@ end
 function transport_from_std_with_rest(::Type{S}, μ::Bind, z::AbstractVector) where {S<:StdMeasure}
     a, z2 = transport_from_std_with_rest(S, μ.α, z)
     b, z_rest = transport_from_std_with_rest(S, _get_β_a(μ, a), z2)
-    return μ.f_c(a, b), z_rest
+    return _combine_variates(μ.f_c, a, b), z_rest
 end
 
 function transport_from_std(::Type{S}, μ::Bind, z::AbstractVector) where {S<:StdMeasure}
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index 0187ba22..fdfe07bd 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -245,7 +245,17 @@ function logdensityof_with_rest(μ::CombinedMeasure{typeof(merge)}, x::NamedTupl
 end
 
 
-rand_impl(ctx::GenContext, μ::CombinedMeasure) = μ.f_c(rand_impl(ctx, μ.α), rand_impl(ctx, μ.β))
+rand_impl(ctx::GenContext, μ::CombinedMeasure) = _combine_variates(μ.f_c, rand_impl(ctx, μ.α), rand_impl(ctx, μ.β))
+
+# Variates of vcat-combined measures are flat streams, nested variates of
+# the components (e.g. of powers of measures with value-dependent sizes)
+# are flattened:
+@inline _combine_variates(f_c, a, b) = f_c(a, b)
+@inline _combine_variates(::typeof(vcat), a, b) = vcat(_flat_stream(a), _flat_stream(b))
+@inline _flat_stream(x::Number) = x
+@inline _flat_stream(x::AbstractArray{<:Number}) = vec(x)
+@inline _flat_stream(x::AbstractArray) = reduce(vcat, map(_flat_stream, x))
+@inline _flat_stream(x::Union{Tuple,NamedTuple}) = reduce(vcat, map(_flat_stream, values(x)))
 
 batched_rand_impl(ctx::GenContext, μ::CombinedMeasure, sz::Dims) = _batched_rand_pointwise(ctx, μ, sz)
 
@@ -256,16 +266,14 @@ end
 
 # Batches of vcat-combined measures are concatenated along the streams:
 function batched_rand_impl(ctx::GenContext, μ::CombinedMeasure{typeof(vcat)}, sz::Dims)
-    _combined_batched_rand(ctx, μ, sz, _static_ndims(μ.α), _static_ndims(μ.β))
+    _combined_batched_rand(ctx, μ, sz, fixed_stream_size(μ))
 end
-function _combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::Dims, k_a::StaticInteger, k_b::StaticInteger)
-    A = _as_stream_batch(batched_rand_impl(ctx, μ.α, sz), k_a)
-    B = _as_stream_batch(batched_rand_impl(ctx, μ.β, sz), k_b)
+function _combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::Dims, ::True)
+    A = _as_stream_batch(batched_rand_impl(ctx, μ.α, sz), μ.α)
+    B = _as_stream_batch(batched_rand_impl(ctx, μ.β, sz), μ.β)
     return vcat(A, B)
 end
-function _combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::Dims, ::Any, ::Any)
-    _batched_rand_pointwise(ctx, μ, sz)
-end
+_combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::Dims, ::False) = _batched_rand_pointwise(ctx, μ, sz)
 
 
 # Transport consumes the variate parts of both component measures in a
@@ -302,7 +310,7 @@ end
 function transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure, z::AbstractVector) where {S<:StdMeasure}
     a, z2 = transport_from_std_with_rest(S, μ.α, z)
     b, z_rest = transport_from_std_with_rest(S, μ.β, z2)
-    return μ.f_c(a, b), z_rest
+    return _combine_variates(μ.f_c, a, b), z_rest
 end
 
 function transport_from_std(::Type{S}, μ::CombinedMeasure, z::AbstractVector) where {S<:StdMeasure}
@@ -372,9 +380,17 @@ function _combined_from_std_with_rest(::Type{S}, μ::CombinedMeasure, z::Abstrac
     transport_from_std_with_rest(S, μ, z)
 end
 function _combined_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, ::Tuple{}) where {S}
+    _combined_batch_from_std_with_rest(S, μ, Z, fixed_stream_size(μ))
+end
+function _combined_batch_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, ::True) where {S}
     A, Z2 = batched_transport_from_std_with_rest(S, μ.α, Z, ())
     B, Z_rest = batched_transport_from_std_with_rest(S, μ.β, Z2, ())
-    return vcat(_as_stream_batch(A, _static_ndims(μ.α)), _as_stream_batch(B, _static_ndims(μ.β))), Z_rest
+    return vcat(_as_stream_batch(A, μ.α), _as_stream_batch(B, μ.β)), Z_rest
+end
+# Components without fixed variate sizes are consumed stream by stream:
+function _combined_batch_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, ::False) where {S}
+    results = map(z -> transport_from_std_with_rest(S, μ, z), sliced(Z, Val(1)))
+    return stacked(map(first, results)), stacked(map(last, results))
 end
 function _combined_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, sz::Dims) where {S}
     _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ))
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index e29aa98a..74f23c56 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -105,7 +105,22 @@ params(d::PowerMeasure) = params(first(marginals(d)))
     basemeasure(d.parent)^d.axes
 end
 
-@inline mspace_ndims(::Type{<:PowerMeasure{M,A}}) where {M,A<:Tuple} = _add_ndims(mspace_ndims(M), fieldcount(A))
+# Numeric flat variates of powers of measures with fixed stream sizes but
+# no variate rank (e.g. tuple products) are streams:
+@inline function mspace_ndims(::Type{<:PowerMeasure{M,A}}) where {M,A<:Tuple}
+    _pwr_ndims(mspace_ndims(M), fieldcount(A), fixed_stream_size(M))
+end
+@inline _pwr_ndims(n::Integer, k::Integer, ::Any) = n + k
+@inline _pwr_ndims(::NoMSpaceElementSize, ::Integer, ::True) = 1
+@inline _pwr_ndims(n::NoMSpaceElementSize, ::Integer, ::False) = n
+
+# Local measures of powers at nested variates are products of the local
+# measures of the elements:
+@inline localmeasure(μ::PowerMeasure, ::AbstractArray{<:Number}) = μ
+function localmeasure(μ::PowerMeasure, x::AbstractArray)
+    size(x) == _dynamic_dims(pwr_size(μ)) || return μ
+    productmeasure(map(Base.Fix1(localmeasure, pwr_base(μ)), x))
+end
 @inline fixed_stream_size(::Type{<:PowerMeasure{M}}) where {M} = fixed_stream_size(M)
 
 # The innermost base measure of nested powers and the total number of power
@@ -120,8 +135,22 @@ end
 # sums the leading dimensions of the result that belong to its axes.
 @inline function _powered_kernel(f::F, μ::PowerMeasure, X) where {F}
     _check_pwr_batch(X, μ)
+    _powered_kernel_impl(f, μ, X, _static_ndims(pwr_base(μ)))
+end
+@inline function _powered_kernel_impl(f::F, μ::PowerMeasure, X, ::Any) where {F}
     _sum_leading_dims(_batched_kernel(f, pwr_base(μ), X), static(length(pwr_axes(μ))))
 end
+# Numeric batches of powers of bases without a variate rank are batches of
+# streams:
+@inline function _powered_kernel_impl(::typeof(logdensityof_impl), μ::PowerMeasure, X::AbstractArray{<:Number}, ::NoMSpaceElementSize)
+    _powered_stream_kernel(μ, X, fixed_stream_size(pwr_base(μ)))
+end
+function _powered_stream_kernel(μ::PowerMeasure, X::AbstractArray, ::True)
+    ℓ, X_rest = batched_logdensityof_with_rest(μ, X, ())
+    size(X_rest, 1) == 0 || _throw_stream_too_long()
+    return ℓ
+end
+_powered_stream_kernel(μ::PowerMeasure, X::AbstractArray, ::False) = _streamwise_ld(logdensityof_impl, μ, X)
 
 # Flat batches of powers have the power dimensions after the variate
 # dimensions of the base measure (where the rank of the base is known):
@@ -172,15 +201,31 @@ end
 end
 
 # Streams: a power consumes its size times the variates of the base measure
-# and sums the base results over its axes.
+# and sums the base results over its axes. Bases without fixed variate
+# sizes are consumed element by element, for single streams.
 function batched_logdensityof_with_rest(μ::PowerMeasure, X::AbstractArray, sz::Dims)
-    _powered_ld_with_rest(μ, X, sz)
+    _powered_ld_with_rest(μ, X, sz, fixed_stream_size(pwr_base(μ)))
+end
+function batched_logdensityof_with_rest(μ::PowerMeasure, x::AbstractVector, sz::Tuple{})
+    _powered_ld_with_rest(μ, x, sz, fixed_stream_size(pwr_base(μ)))
 end
-batched_logdensityof_with_rest(μ::PowerMeasure, x::AbstractVector, sz::Tuple{}) = _powered_ld_with_rest(μ, x, sz)
-function _powered_ld_with_rest(μ::PowerMeasure, X::AbstractArray, sz::Dims)
-    ℓ, X_rest = batched_logdensityof_with_rest(pwr_base(μ), X, (map(dynamic, _size_dims(pwr_size(μ)))..., sz...))
+function _powered_ld_with_rest(μ::PowerMeasure, X::AbstractArray, sz::Dims, ::True)
+    ℓ, X_rest = batched_logdensityof_with_rest(pwr_base(μ), X, (_dynamic_dims(pwr_size(μ))..., sz...))
     return _sum_leading_dims(ℓ, static(length(pwr_axes(μ)))), X_rest
 end
+function _powered_ld_with_rest(μ::PowerMeasure, x::AbstractVector, ::Tuple{}, ::False)
+    ν = pwr_base(μ)
+    ℓ = zero(_logd_numtype(x))
+    x_rest = x
+    for _ in 1:length(marginals(μ))
+        ℓ_i, _, x_rest = logdensityof_with_rest(ν, x_rest)
+        ℓ += ℓ_i
+    end
+    return ℓ, x_rest
+end
+@noinline function _powered_ld_with_rest(μ::PowerMeasure, ::AbstractArray, ::Dims, ::False)
+    throw(ArgumentError("Batches of variate streams containing powers of measures of type $(nameof(typeof(pwr_base(μ)))) must be consumed stream by stream"))
+end
 
 # Support checks of powers run over the flat variate storage where the
 # innermost base measure has scalar variates, elementwise otherwise:
@@ -193,6 +238,13 @@ end
     ν, _ = _pwr_unwrap(μ)
     _all_insupport(broadcast(_insupport_bool ∘ Base.Fix1(insupport, ν), x_flat))
 end
+@inline function _powered_insupport(μ::PowerMeasure, x, x_flat::AbstractArray, ::StaticInteger{K}) where {K}
+    ν, _ = _pwr_unwrap(μ)
+    _all_insupport(map(_insupport_bool ∘ Base.Fix1(insupport, ν), sliced(x_flat, Val(K))))
+end
+@inline function _powered_insupport(μ::PowerMeasure, x, ::AbstractArray{<:Number}, ::NoMSpaceElementSize)
+    NoFastInsupport{typeof(μ)}()
+end
 @inline _powered_insupport(μ::PowerMeasure, x, ::Any, ::Any) = _powered_insupport_elementwise(pwr_base(μ), x)
 
 @inline function _powered_insupport_elementwise(ν, x::AbstractArray)
@@ -241,9 +293,22 @@ massof(m::PowerMeasure) = massof(m.parent)^dynamic(size2length(pwr_size(m)))
 
 function batched_transport_to_std(::Type{S}, μ::PowerMeasure, X::AbstractArray) where {S<:StdMeasure}
     _check_pwr_batch(X, μ)
+    _pwr_batched_to_std(S, μ, X, _static_ndims(pwr_base(μ)))
+end
+function batched_transport_to_std(::Type{S}, μ::PowerMeasure, X::Union{Tuple,NamedTuple}) where {S<:StdMeasure}
+    _pwr_batched_to_std(S, μ, X, nothing)
+end
+@inline function _pwr_batched_to_std(::Type{S}, μ::PowerMeasure, X, ::Any) where {S}
     ν, n = _pwr_unwrap(μ)
     _merge_leading_dims(batched_transport_to_std(S, ν, X), static(1) + n)
 end
+# Numeric batches of powers of bases without a variate rank are batches of
+# streams:
+function _pwr_batched_to_std(::Type{S}, μ::PowerMeasure, X::AbstractArray{<:Number}, ::NoMSpaceElementSize) where {S}
+    Z, X_rest = batched_transport_to_std_with_rest(S, μ, X, ())
+    size(X_rest, 1) == 0 || _throw_stream_too_long()
+    return Z
+end
 
 function batched_transport_from_std(::Type{S}, μ::PowerMeasure, Z::AbstractArray) where {S<:StdMeasure}
     ν, _ = _pwr_unwrap(μ)
@@ -299,9 +364,19 @@ function transport_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVect
     _pwr_point_to_std_with_rest(S, μ, x, fixed_stream_size(pwr_base(μ)))
 end
 function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::True) where {S}
+    _pwr_point_to_std_with_rest(S, μ, x, _static_ndims(pwr_base(μ)))
+end
+function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::StaticInteger) where {S}
     x_μ, x_rest = _consume_from_stream(x, _stream_consume_size(μ))
     return _as_stdstream(transport_to_std(S, μ, x_μ)), x_μ, x_rest
 end
+# Bases without a variate rank (tuple products) consume streams via the
+# batched protocol:
+function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::NoMSpaceElementSize) where {S}
+    z, x_rest = _pwr_to_std_with_rest(S, μ, x, (), static(true))
+    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    return z, x_μ, x_rest
+end
 function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::False) where {S}
     ν = pwr_base(μ)
     zs = Vector{Any}(undef, length(marginals(μ)))
@@ -323,6 +398,11 @@ function _pwr_from_std_with_rest(::Type{S}, μ, z, ::AbstractNoDOF) where {S}
     _marginals_from_std_with_rest(S, marginals(μ), z)
 end
 
+# The stream length of a power with a base of fixed stream length:
+@inline function _fixed_stream_length(μ::PowerMeasure)
+    _fixed_stream_length(pwr_base(μ)) * prod(_dynamic_dims(pwr_size(μ)))
+end
+
 # The nested variate layout of a power over its flat storage, batches of
 # tuple variates become struct arrays:
 @inline _pwr_variate(μ::PowerMeasure, A::AbstractArray) = _pwr_variate_impl(μ, A)
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index b134e067..3669f358 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -67,6 +67,14 @@ end
 proxy(μ::ProductMeasure{<:FillArrays.Fill}) =
     powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ)))
 
+# Array products with fused kernels draw their variates in one batch (also
+# on devices):
+function rand_impl(ctx::GenContext, d::ProductMeasure{<:AbstractArray{M}}) where {M}
+    _array_product_rand(ctx, d, _fused_marginals(M))
+end
+_array_product_rand(ctx::GenContext, d::ProductMeasure, ::Val{true}) = batched_rand_impl(ctx, d, ())
+_array_product_rand(ctx::GenContext, d::ProductMeasure, ::Val{false}) = _map(Base.Fix1(_marginal_rand, ctx), marginals(d))
+
 # Batches of tuple and named tuple products are tuples resp. named tuples
 # of marginal batches:
 function batched_rand_impl(ctx::GenContext, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, sz::Dims)
@@ -281,8 +289,9 @@ end
 # TODO: Better `map` support in MappedArrays
 _map(f, args...) = map(f, args...)
 _map(f, x::MappedArrays.ReadonlyMappedArray) = mappedarray(fchain((x.f, f)), x.data)
-# Variates of struct array marginals are collected into plain arrays:
-_map(f, x::StructArray) = map(f, collect(x))
+# `map` over a struct array builds struct arrays of the results, variates
+# of the marginals are wanted as plain arrays:
+_map(f, x::StructArray) = [f(m) for m in x]
 
 function testvalue(::Type{T}, d::AbstractProductMeasure) where {T}
     _map(m -> testvalue(T, m), marginals(d))
@@ -407,6 +416,34 @@ function transport_from_std(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTup
     return x
 end
 
+# Streams of tuple product variates are consumed marginal by marginal:
+function transport_to_std_with_rest(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, x::AbstractVector) where {S<:StdMeasure}
+    z, x_rest = _marginals_to_std_with_rest(S, values(marginals(μ)), x)
+    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    return z, x_μ, x_rest
+end
+
+function batched_transport_to_std_with_rest(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, X::AbstractArray, sz::Dims) where {S<:StdMeasure}
+    _tuple_product_to_std_with_rest(S, μ, X, sz)
+end
+function _tuple_product_to_std_with_rest(::Type{S}, μ, X::AbstractArray, ::Tuple{}) where {S}
+    _marginals_to_std_with_rest(S, values(marginals(μ)), X)
+end
+function _tuple_product_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::Dims) where {S}
+    X_v, X_rest = _split_stream_variates(μ, X, sz)
+    Z, _ = _marginals_to_std_with_rest(S, values(marginals(μ)), X_v)
+    return _merge_multiplicity(Z, sz), X_rest
+end
+
+function _marginals_to_std_with_rest(::Type{S}, νs::Tuple{Vararg{Any}}, X::AbstractArray) where {S}
+    Z1, X_rest = batched_transport_to_std_with_rest(S, νs[1], X, ())
+    Z2_end, X_final_rest = _marginals_to_std_with_rest(S, Base.tail(νs), X_rest)
+    return vcat(Z1, Z2_end), X_final_rest
+end
+function _marginals_to_std_with_rest(::Type{S}, νs::Tuple{Any}, X::AbstractArray) where {S}
+    batched_transport_to_std_with_rest(S, νs[1], X, ())
+end
+
 function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:Tuple}, z::AbstractVector) where {S<:StdMeasure}
     _marginals_from_std_with_rest(S, marginals(μ), z)
 end
@@ -608,12 +645,33 @@ end
 
 # Streams: tuple products consume marginal by marginal, so marginals of
 # value-dependent size are supported for a single variate per stream.
+# Several variates per stream are split by the fixed stream length of the
+# product.
 function batched_logdensityof_with_rest(μ::ProductMeasure{<:Tuple}, X::AbstractArray, ::Tuple{})
     _marginals_ld_with_rest(marginals(μ), X)
 end
 function batched_logdensityof_with_rest(μ::ProductMeasure{<:Tuple}, x::AbstractVector, ::Tuple{})
     _marginals_ld_with_rest(marginals(μ), x)
 end
+function batched_logdensityof_with_rest(μ::ProductMeasure{<:Tuple}, X::AbstractArray, sz::Dims)
+    X_v, X_rest = _split_stream_variates(μ, X, sz)
+    ℓ, _ = _marginals_ld_with_rest(marginals(μ), X_v)
+    return ℓ, X_rest
+end
+function batched_logdensityof_with_rest(μ::ProductMeasure{<:NamedTuple{names}}, X::AbstractArray, sz::Dims) where {names}
+    batched_logdensityof_with_rest(productmeasure(values(marginals(μ))), X, sz)
+end
+
+# The rows of `prod(sz)` variates of fixed stream length, as a batch of
+# streams `(stream length, sz..., batch dims...)`:
+function _split_stream_variates(μ, X::AbstractArray, sz::Dims)
+    n_rows = _fixed_stream_length(μ)
+    X_μ, X_rest = _batched_split(X, n_rows * prod(sz))
+    return reshape(X_μ, (n_rows, sz..., Base.tail(size(X_μ))...)), X_rest
+end
+
+@inline _fixed_stream_length(μ::ProductMeasure{<:Tuple}) = sum(_fixed_stream_length, marginals(μ))
+@inline _fixed_stream_length(μ::ProductMeasure{<:NamedTuple}) = sum(_fixed_stream_length, values(marginals(μ)))
 function _marginals_ld_with_rest(ms::Tuple, X::AbstractArray)
     ℓ1, X2 = batched_logdensityof_with_rest(ms[1], X, ())
     ℓ_rest, X_rest = _marginals_ld_with_rest(Base.tail(ms), X2)
diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
index e6ba3867..8a402036 100644
--- a/src/combinators/smart-constructors.jl
+++ b/src/combinators/smart-constructors.jl
@@ -42,7 +42,7 @@ end
 @inline _generic_powermeasure_stage2(μ::AbstractMeasure, exponent::Tuple) =
     PowerMeasure(μ, exponent)
 
-@inline function _generic_powermeasure_stage2(μ::Dirac, exponent::Tuple)
+@inline function _generic_powermeasure_stage2(μ::Dirac{<:Number}, exponent::Tuple)
     Dirac(maybestatic_fill(μ.x, exponent))
 end
 
@@ -65,7 +65,7 @@ Examples:
 ```julia
 productmeasure((StdNormal(), StdExponential()))
 productmeasure((a = StdNormal(), b = StdExponential()))
-productmeasure([pushfwd(Base.Fix1(*, scale), StdExponential()) for scale in 0.1:0.2:2])
+productmeasure([pushfwd(AffineMaps.Mul(scale), StdExponential()) for scale in 0.1:0.2:2])
 ```
 """
 function productmeasure end
@@ -135,9 +135,12 @@ function _marginal_storage(mar::AbstractArray{T}) where {T}
 end
 _marginal_storage(mar::StructArray) = mar
 
+# Function objects stay opaque columns, their wrappers (e.g. `Base.Fix1`)
+# can't be rebuilt from their fields via ConstructionBase:
 @inline function _unwrap_field(::Type{T}) where {T}
     isstructtype(T) && !Base.issingletontype(T) && fieldcount(T) > 0 &&
-        !(T <: Number) && !(T <: AbstractArray) && !(T <: Tuple) && !(T <: AbstractString) && !(T <: Symbol)
+        !(T <: Number) && !(T <: AbstractArray) && !(T <: Tuple) && !(T <: AbstractString) &&
+        !(T <: Symbol) && !(T <: Function)
 end
 
 @inline function _generic_productmeasure_impl(
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index 7fffad77..948adf05 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -30,7 +30,7 @@ end
     _lazy_add(_logweight_for(d.logweight, X), batched_logdensityof_impl(basemeasure(d), X))
 end
 @inline function batched_logdensity_def(d::AbstractWeightedMeasure, X)
-    _lazy_add(_logweight_for(d.logweight, X), _zero_logd_batch(X, mspace_ndims(basemeasure(d))))
+    _lazy_add(_logweight_for(d.logweight, X), _zero_logd_batch(X, _static_ndims(basemeasure(d))))
 end
 
 @inline rand_impl(ctx::GenContext, μ::AbstractWeightedMeasure) = rand_impl(ctx, basemeasure(μ))
diff --git a/src/density-batched.jl b/src/density-batched.jl
index ea6e1f99..91d9f6df 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -23,10 +23,10 @@ or an array of numbers for measures with scalar variates), the flat storage
 of a batch (an array of numbers with the variate dimensions leading, see
 [`MeasureBase.mspace_ndims`](@ref)), or a tuple resp. `NamedTuple` of
 batches for measures with tuple resp. `NamedTuple` variates. Returns an
-array over the batch dimensions, semantically equivalent to
-`logdensityof.(Ref(μ), X)` for arrays of variates. Batches with flat
-storage are evaluated in fused operations, compatible with GPU and traced
-arrays.
+array over the batch dimensions (a number for a single variate),
+semantically equivalent to `logdensityof.(Ref(μ), X)` for arrays of
+variates. Batches with flat storage are evaluated in fused operations,
+compatible with GPU and traced arrays.
 
 Measure types implement [`MeasureBase.batched_logdensityof_impl`](@ref).
 """
@@ -126,6 +126,8 @@ struct NoFlatStorage end
 end
 @inline _flat_storage_bymode(::AbstractArray, ::UnknownSplitMode) = NoFlatStorage()
 @inline _flat_storage_bymode(::AbstractArray, ::NonSplitMode) = NoFlatStorage()
+# Ragged batches (e.g. `VectorOfArrays`) are evaluated variate by variate:
+@inline _flat_storage_bymode(::AbstractArray, ::AbstractPartMode) = NoFlatStorage()
 
 @inline _flat_storage(X::StructArray{<:Union{Tuple,NamedTuple}}) = _components_storage(StructArrays.components(X))
 @inline function _flat_storage(X::AbstractArray{<:Union{Tuple,NamedTuple}})
@@ -173,13 +175,16 @@ end
     throw(ArgumentError("Batched density kernel returned a result of size $sz_result for a batch of size $sz_batch, the variate dimensions of the measure don't match the batch"))
 end
 
-# Point evaluation: array variates go through the batched kernel with zero
-# batch dimensions, other variates through the point kernel. A batched
-# kernel that returns an array for a single variate has taken the variate
-# for a batch: the variate doesn't fit the measure, or the measure lacks a
-# batched kernel for array variates.
-@inline _point_ld(f::F, μ, x::AbstractArray{<:Number}) where {F} = _point_result(_materialize(_batched_kernel(f, μ, x)), μ)
+# Point evaluation: array variates of measures with a declared variate rank
+# go through the batched kernel with zero batch dimensions, other variates
+# through the point kernel. A batched kernel that returns an array for a
+# single variate has taken the variate for a batch: the variate doesn't
+# fit the measure, or the measure lacks a batched kernel for array
+# variates.
+@inline _point_ld(f::F, μ, x::AbstractArray{<:Number}) where {F} = _point_ld_byrank(f, μ, x, _static_ndims(μ))
 @inline _point_ld(f::F, μ, x) where {F} = f(μ, x)
+@inline _point_ld_byrank(f::F, μ, x, ::StaticInteger) where {F} = _point_result(_materialize(_batched_kernel(f, μ, x)), μ)
+@inline _point_ld_byrank(f::F, μ, x, ::NoMSpaceElementSize) where {F} = f(μ, x)
 @inline _point_ld(f::F, μ::PrimitiveMeasure, x::AbstractArray{<:Number}) where {F} = f(μ, x)
 
 @inline _point_result(ℓ::Number, μ) = ℓ
@@ -220,7 +225,7 @@ end
 @inline _sum_leading_dims_lazy(bc::_LazyBroadcast, ::StaticInteger{0}, ::StaticInteger{0}) = bc
 @inline _sum_leading_dims_lazy(bc::_EagerReducibleBroadcast, ::StaticInteger{0}, ::StaticInteger{0}) = bc
 @inline function _sum_leading_dims_lazy(bc::_EagerReducibleBroadcast, ::StaticInteger{N}, ::StaticInteger{N}) where {N}
-    isempty(bc) ? sum(copy(bc)) : sum(bc)
+    length(bc) == 0 ? sum(copy(bc)) : sum(bc)
 end
 @inline _sum_leading_dims_lazy(bc::_LazyBroadcast, ::StaticInteger{N}, ::StaticInteger{N}) where {N} = sum(copy(bc))
 @inline function _sum_leading_dims_lazy(bc::_LazyBroadcast, n::StaticInteger, ::StaticInteger)
@@ -247,6 +252,9 @@ end
 end
 @inline _zero_logd_batch(X::AbstractArray{<:Any,N}, ::StaticInteger{N}) where {N} = zero(_logd_numtype(X))
 @inline _zero_logd_batch(x::Number, ::StaticInteger{0}) = zero(_logd_numtype(x))
+@noinline function _zero_logd_batch(X, ::NoMSpaceElementSize{MU}) where {MU}
+    throw(ArgumentError("Batched density evaluation for measures of type $(nameof(MU)) requires MeasureBase.mspace_ndims to be declared for the type or MeasureBase.batched_logdensity_def to be implemented"))
+end
 
 
 # Streams: variates of composed measures are consumed from flat vector
@@ -268,7 +276,9 @@ kernel evaluation, the default implementation does so for the variate size
 given by [`MeasureBase.mspace_flatsize`](@ref) or
 [`MeasureBase.some_mspace_elsize`](@ref). Measure types with variates of
 value-dependent size implement `batched_logdensityof_with_rest` for
-`sz == ()` themselves, they can only be evaluated stream by stream.
+single streams and `sz == ()` themselves and report
+`MeasureBase.fixed_stream_size` as false, so that the enclosing stream
+combinators consume batches stream by stream.
 """
 function batched_logdensityof_with_rest end
 
@@ -288,28 +298,39 @@ function _stream_ld_with_rest(f::F, μ, X::AbstractArray, sz::Dims) where {F}
     return _consumed_ld(f, μ, X_μ, vsz), X_rest
 end
 
-# Scalar variates are consumed as `(1, sz..., batch dims...)` and the leading
-# dimension is summed out, so that no rank-0 arrays arise:
-@inline _consumed_ld(f::F, μ, X_μ, ::Tuple{}) where {F} = _sum_leading_dims(_batched_kernel(f, μ, X_μ), static(1))
+# Scalar variates are consumed as `(1, sz..., batch dims...)`, their
+# leading dimension is dropped before the kernel runs:
+@inline _consumed_ld(f::F, μ, X_μ, ::Tuple{}) where {F} = _batched_kernel(f, μ, _drop_stdstream_dim(X_μ))
 @inline _consumed_ld(f::F, μ, X_μ, ::SizeLike) where {F} = _batched_kernel(f, μ, X_μ)
 
 # Consume `prod(sz)` variates of flat size `vsz` from the leading rows of a
 # batch of streams as a flat batch `(vsz..., sz..., batch dims...)`; scalar
-# variates as `(1, sz..., batch dims...)`:
+# variates as `(1, sz..., batch dims...)`. Static sizes keep static
+# streams static.
 @inline function _batched_consume(X::AbstractArray, vsz::SizeLike, sz::Dims)
     dims = _consumed_dims(vsz)
-    n_rows = prod(dims) * prod(sz)
-    X_flat, X_rest = _batched_split(X, n_rows)
+    X_flat, X_rest = _batched_split(X, _chunk_rows(prod(dims), sz))
     return _reshape_consumed(X_flat, (dims..., sz...)), X_rest
 end
-@inline _consumed_dims(::Tuple{}) = (1,)
-@inline _consumed_dims(vsz::SizeLike) = map(dynamic, _size_dims(vsz))
+@inline _consumed_dims(::Tuple{}) = (static(1),)
+@inline _consumed_dims(vsz::SizeLike) = _size_dims(vsz)
+@inline _chunk_rows(n::IntegerLike, ::Tuple{}) = n
+@inline _chunk_rows(n::IntegerLike, sz::Dims) = dynamic(n) * prod(sz)
 
-@inline _reshape_consumed(X_flat::AbstractArray, ::Tuple{Int}) = X_flat
-@inline function _reshape_consumed(X_flat::AbstractArray, dims::Tuple{Vararg{Int}})
-    reshape(X_flat, (dims..., Base.tail(size(X_flat))...))
+@inline _reshape_consumed(X_flat::AbstractArray, ::Tuple{IntegerLike}) = X_flat
+@inline function _reshape_consumed(X_flat::AbstractArray, dims::Tuple{Vararg{IntegerLike}})
+    _reshape_batch(X_flat, (dims..., Base.tail(_batch_dims(X_flat))...))
 end
 
+# Sizes as tuples of (maybe static) integers, reshapes that keep static
+# arrays static, and the leading dimension of a batch of streams:
+@inline _batch_dims(A::AbstractArray) = _size_dims(maybestatic_size(A))
+@inline _reshape_batch(A::AbstractArray, dims::Tuple) = reshape(A, map(dynamic, dims))
+@inline _reshape_batch(A::StaticArray, dims::Tuple{Vararg{StaticInteger}}) = maybestatic_reshape(A, dims)
+@inline _as_stdstream_batch(Z::AbstractArray) = _reshape_batch(Z, (static(1), _batch_dims(Z)...))
+@inline _as_stdstream_batch(z::Number) = SVector(z)
+@inline _drop_stdstream_dim(Z::AbstractArray) = _reshape_batch(Z, Base.tail(_batch_dims(Z)))
+
 @inline function _batched_split(A::AbstractArray, n::IntegerLike)
     n_rows = dynamic(n)
     stream_idxs = axes(A, 1)
@@ -323,7 +344,7 @@ end
     return A_flat, A_rest
 end
 
-# Static streams split into static chunks:
+# Static streams split into static chunks for static row counts:
 @inline function _batched_split(A::StaticVector, n_rows::StaticInteger{N}) where {N}
     idxs = maybestatic_eachindex(A)
     i0 = maybestatic_first(idxs)
diff --git a/src/getdof.jl b/src/getdof.jl
index c2edd15c..aca22e67 100644
--- a/src/getdof.jl
+++ b/src/getdof.jl
@@ -94,8 +94,12 @@ function some_dof(μ)
 end
 
 _try_direct_dof(::AbstractMeasure, dof::IntegerLike) = dof
-_try_direct_dof(μ::AbstractMeasure, ::AbstractNoDOF) =
-    _try_local_dof(μ, some_dof(_some_localmeasure(μ)))
+function _try_direct_dof(μ::AbstractMeasure, ::AbstractNoDOF)
+    μ_local = _some_localmeasure(μ)
+    # A local measure of the same type would recurse forever:
+    typeof(μ_local) === typeof(μ) && _try_local_dof(μ, NoDOF{typeof(μ)}())
+    _try_local_dof(μ, some_dof(μ_local))
+end
 
 _try_local_dof(::AbstractMeasure, dof::IntegerLike) = dof
 _try_local_dof(μ::AbstractMeasure, ::AbstractNoDOF) =
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 0bf69750..5a92d639 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -58,13 +58,17 @@ end
 @inline transport_from_std(::Type{S}, μ::Dirac, z::AbstractVector) where {S<:StdMeasure} = μ.x
 @inline transport_from_std_with_rest(::Type{S}, μ::Dirac, z::AbstractVector) where {S<:StdMeasure} = μ.x, z
 
+# Batched kernels cover Dirac measures with numbers and numeric arrays as
+# flat variates, others have no declared variate rank:
+const _FlatDirac = Dirac{<:Union{Number,AbstractArray{<:Number}}}
+
 @inline batched_transport_to_std(::Type{S}, ::Dirac, ::Number) where {S<:StdMeasure} = SVector{0,Bool}()
-function batched_transport_to_std(::Type{S}, μ::Dirac, X::AbstractArray) where {S<:StdMeasure}
+function batched_transport_to_std(::Type{S}, μ::_FlatDirac, X::AbstractArray) where {S<:StdMeasure}
     n = length(_value_flatsize(μ.x))
     similar(X, Bool, (0, ntuple(i -> size(X, n + i), Val(ndims(X) - n))...))
 end
 
-function batched_transport_from_std(::Type{S}, μ::Dirac, Z::AbstractArray) where {S<:StdMeasure}
+function batched_transport_from_std(::Type{S}, μ::_FlatDirac, Z::AbstractArray) where {S<:StdMeasure}
     _const_variates(μ.x, Z)
 end
 @inline _const_variates(x::Number, ::AbstractVector) = x
@@ -74,14 +78,16 @@ function _const_variates(x, Z::AbstractArray)
     return X
 end
 
-@inline mspace_ndims(::Type{<:Dirac{<:AbstractArray{<:Any,N}}}) where {N} = N
+@inline mspace_ndims(::Type{<:Dirac{<:AbstractArray{<:Number,N}}}) where {N} = N
 
 # Batches of array variates: all elements of a variate must match.
-function batched_logdensityof_impl(μ::Dirac{<:AbstractArray{<:Any,N}}, X::AbstractArray) where {N}
+function batched_logdensityof_impl(μ::Dirac{<:AbstractArray{<:Number,N}}, X::AbstractArray) where {N}
     matches = _all_leading_dims(X .== μ.x, static(N))
     ifelse.(matches, zero(_logd_numtype(X)), _neg_inf_logd(X))
 end
-batched_logdensity_def(μ::Dirac{<:AbstractArray}, X::AbstractArray) = _zero_logd_batch(X, static(ndims(μ.x)))
+function batched_logdensity_def(μ::Dirac{<:AbstractArray{<:Number}}, X::AbstractArray)
+    _zero_logd_batch(X, static(ndims(μ.x)))
+end
 
 @inline _all_leading_dims(A::AbstractArray{Bool,N}, ::StaticInteger{N}) where {N} = all(A)
 @inline function _all_leading_dims(A::AbstractArray{Bool}, ::StaticInteger{N}) where {N}
diff --git a/src/transport-batched.jl b/src/transport-batched.jl
index 40bd9c56..4b476bfb 100644
--- a/src/transport-batched.jl
+++ b/src/transport-batched.jl
@@ -45,18 +45,6 @@ end
     stacked(map(Base.Fix1(_ToStd{S}(), μ), sliced(X, Val(K))))
 end
 
-# Standard variates of scalar-variate measures form the first dimension of
-# a batch of streams:
-@inline _as_stdstream_batch(Z::AbstractArray) = _reshape_batch(Z, (static(1), _batch_dims(Z)...))
-@inline _as_stdstream_batch(z::Number) = SVector(z)
-@inline _drop_stdstream_dim(Z::AbstractArray) = _reshape_batch(Z, Base.tail(_batch_dims(Z)))
-
-# Sizes as tuples of (maybe static) integers, and reshapes that keep static
-# arrays static:
-@inline _batch_dims(A::AbstractArray) = _size_dims(maybestatic_size(A))
-@inline _reshape_batch(A::AbstractArray, dims::Tuple) = reshape(A, map(dynamic, dims))
-@inline _reshape_batch(A::StaticArray, dims::Tuple{Vararg{StaticInteger}}) = maybestatic_reshape(A, dims)
-
 # Merge the leading `N` dimensions of an array into one, `N == 0` adds a
 # leading dimension of size one:
 @inline function _merge_leading_dims(A::AbstractArray, ::StaticInteger{N}) where {N}
@@ -67,13 +55,29 @@ end
 end
 @inline _merge_leading_dims(A::AbstractArray, ::StaticInteger{0}) = _reshape_batch(A, (static(1), _batch_dims(A)...))
 
-# A flat batch of variates as a batch of streams, the variate dimensions
-# merged into the first dimension:
+# A flat batch of variates of `μ` as a batch of streams, the variate
+# dimensions merged into the first dimension. Tuples of batches (tuple
+# products and their powers) interleave the rows of their components
+# variate by variate.
+@inline _as_stream_batch(X, μ) = _as_stream_batch(X, _static_ndims(μ))
 @inline _as_stream_batch(X::AbstractArray, ::StaticInteger{K}) where {K} = _merge_leading_dims(X, static(K))
 @inline _as_stream_batch(x::Number, ::StaticInteger{0}) = SVector(x)
 @noinline function _as_stream_batch(X, ::NoMSpaceElementSize)
     throw(ArgumentError("Concatenating batches of variates requires MeasureBase.mspace_ndims to be declared for the measures involved"))
 end
+@inline function _as_stream_batch(X::Union{Tuple,NamedTuple}, μ::ProductMeasure)
+    vcat(map(_as_stream_batch, values(X), values(marginals(μ)))...)
+end
+function _as_stream_batch(X::Union{Tuple,NamedTuple}, μ::PowerMeasure)
+    ν, _ = _pwr_unwrap(μ)
+    n_pwr = length(_pwr_dims(μ))
+    n = prod(map(dynamic, _pwr_dims(μ)))
+    parts = map(values(X), values(marginals(ν))) do Xi, m
+        A = _as_stream_batch(Xi, m)
+        reshape(A, (size(A, 1), n, ntuple(i -> size(A, 1 + n_pwr + i), Val(ndims(A) - 1 - n_pwr))...))
+    end
+    _merge_leading_dims(vcat(parts...), static(2))
+end
 
 # The standard variates of a single variate must form a vector:
 @inline _single_std(z::AbstractVector) = z
@@ -175,8 +179,8 @@ dimension along the streams) and transport them to `μ`.
 
 Returns a tuple `(X, Z_rest)` of the flat batch `(flat variate dims...,
 sz..., batch dims...)` of variates of `μ` and the unconsumed rest of the
-streams. The default implementation consumes [`getdof(μ)`](@ref) entries
-per variate, a single stream with `sz == ()` goes through
+streams. The default implementation consumes [`MeasureBase.fast_dof(μ)`](@ref)
+entries per variate, a single stream with `sz == ()` goes through
 [`MeasureBase.transport_from_std_with_rest`](@ref). Measures whose
 variates are composed of the variates of other measures implement
 `batched_transport_from_std_with_rest` in terms of their components.
@@ -199,8 +203,6 @@ end
 @noinline function _batched_from_std_bydof(::Type{S}, μ, ::AbstractArray, ::Dims, ::AbstractNoDOF) where {S}
     throw(ArgumentError("Batched transport from standard measures requires measures of type $(nameof(typeof(μ))) to have fast degrees of freedom or to implement MeasureBase.batched_transport_from_std_with_rest"))
 end
-@inline _chunk_rows(n::IntegerLike, ::Tuple{}) = n
-@inline _chunk_rows(n::IntegerLike, sz::Dims) = dynamic(n) * prod(sz)
 
 
 """
@@ -236,7 +238,12 @@ function Broadcast.broadcasted(f::TransportFunction, bc::Broadcast.Broadcasted)
     Broadcast.broadcasted(f, Broadcast.materialize(bc))
 end
 
-Broadcast.broadcasted(f::TransportFunction, X::StaticArray) = map(_Pointwise(f), X)
+# Static arrays of scalar variates are transported point by point:
+Broadcast.broadcasted(f::TransportFunction, X::StaticArray) = _broadcast_static(f, X, _static_ndims(f.μ))
+_broadcast_static(f::TransportFunction, X::StaticArray, ::StaticInteger{0}) = map(_Pointwise(f), X)
+function _broadcast_static(f::TransportFunction, X::StaticArray, k)
+    _broadcast_transport(f, X, X, k, _static_ndims(f.ν))
+end
 
 function _broadcast_transport(f::TransportFunction, X, X_flat::AbstractArray, ::StaticInteger, ::StaticInteger{K}) where {K}
     Y_flat = batched_transport_def(f.ν, f.μ, X_flat)
@@ -267,7 +274,6 @@ end
 # result, nested powers included, batches of tuple variates as struct
 # arrays:
 @inline _batch_variates(Y::AbstractArray, ν, ::Val{K}) where {K} = _nest_batch(Y, Val(K))
-@inline _batch_variates(Y::Union{Tuple,NamedTuple}, ν, ::Val) = _pwr_variate(ν, Y)
 @inline function _batch_variates(Y::AbstractArray, ν::PowerMeasure, ::Val)
     sliced(_pwr_variate(ν, Y), Val(length(pwr_axes(ν))))
 end
diff --git a/test/batched_regressions.jl b/test/batched_regressions.jl
new file mode 100644
index 00000000..0456dc58
--- /dev/null
+++ b/test/batched_regressions.jl
@@ -0,0 +1,134 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+# Regression tests for the batched-first review findings.
+
+using Test
+
+using MeasureBase
+using MeasureBase: StdNormal, StdUniform, StdExponential, Dirac, Lebesgue, GenContext
+using MeasureBase: productmeasure, pushfwd, mcombine, mbind, weightedmeasure, insupport, testvalue
+using MeasureBase: batched_rand_impl, batched_transport_to_std_with_rest, batched_transport_from_std_with_rest
+using MeasureBase.InverseFunctions: inverse
+using ArraysOfArrays: VectorOfVectors, nestedview, flatview, sliced
+using StaticArrays: SVector
+using Static: static
+using Distributions: MvNormal, logpdf
+using LinearAlgebra: I
+using JLArrays
+
+struct UnknownRankMeasure <: AbstractMeasure end
+MeasureBase.basemeasure(::UnknownRankMeasure) = Lebesgue()
+MeasureBase.logdensity_def(::UnknownRankMeasure, x) = -sum(abs, x)
+MeasureBase.insupport(::UnknownRankMeasure, x) = true
+
+# A parameterized function object (not unwrapped into struct array columns):
+struct Scale <: Function
+    s::Float64
+end
+(f::Scale)(x) = f.s * x
+MeasureBase.InverseFunctions.inverse(f::Scale) = Scale(inv(f.s))
+MeasureBase.ChangesOfVariables.with_logabsdet_jacobian(f::Scale, x) = (f(x), log(abs(f.s)))
+
+@testset "batched regressions" begin
+    @testset "products of function-wrapper marginals" begin
+        pm = productmeasure([pushfwd(Scale(s), StdExponential()) for s in 0.1:0.2:0.9])
+        x = rand(pm)
+        @test logdensityof(pm, x) ≈ sum(logdensityof.(MeasureBase.marginals(pm), x))
+        X = rand(pm^4)
+        @test logdensities(pm, X) ≈ [logdensityof(pm, X[j]) for j in 1:4]
+        f = transport_to(StdUniform()^5, pm)
+        @test flatview(inverse(f).(f.(sliced(flatview(X), Val(1))))) ≈ flatview(X)
+        JLArrays.allowscalar(false)
+        @test Array(logdensities(MeasureBase.Adapt.adapt(JLArray, pm), JLArray(flatview(X)))) ≈ logdensities(pm, X)
+    end
+
+    @testset "support of powers with array-variate bases" begin
+        mv = MeasureBase.AsMeasure{typeof(MvNormal(zeros(2), I(2)))}(MvNormal(zeros(2), I(2)))
+        Xf = randn(2, 3)
+        @test insupport(mv^3, Xf) === true
+        @test insupport(mv^3, nestedview(Xf)) === true
+        @test logdensities(mv, Xf) ≈ [logpdf(mv.obj, Xf[:, j]) for j in 1:3]
+        @test logdensityof(mv^3, Xf) ≈ sum(logpdf(mv.obj, Xf))
+    end
+
+    @testset "powers of array Diracs" begin
+        D = Dirac([1.0, 2.0])
+        μ = D^3
+        x = [1.0 1.0 1.0; 2.0 2.0 2.0]
+        @test MeasureBase.mspace_ndims(typeof(μ)) == 2
+        @test logdensityof(μ, x) == 0 && insupport(μ, x)
+        @test logdensityof(μ, [[1.0, 2.0] for _ in 1:3]) == 0
+        @test logdensityof(μ, 2 .* x) == -Inf
+        @test logdensities(μ, cat(x, 2 .* x; dims = 3)) == [0.0, -Inf]
+        @test rand(μ) == nestedview(x)
+    end
+
+    @testset "streams with value-dependent sizes inside powers" begin
+        bnd = mbind(a -> StdNormal()^(a > 0.5 ? 2 : 1), StdUniform(), vcat)
+        @test testvalue(bnd) isa AbstractVector
+        m = mcombine(vcat, StdNormal(), bnd^2)
+        x = [0.5, 0.7, 0.1, 0.2, 0.3, 0.1]
+        ℓ = logdensityof(m, x)
+        @test ℓ ≈ logdensityof(StdNormal(), 0.5) + logdensityof(bnd, [0.7, 0.1, 0.2]) + logdensityof(bnd, [0.3, 0.1])
+        @test logdensities(m, hcat(x, x)) ≈ [ℓ, ℓ]
+        f = transport_to(StdUniform()^6, m)
+        @test inverse(f)(f(x)) ≈ x
+        @test flatview(inverse(f).(f.(sliced(hcat(x, x), Val(1))))) ≈ hcat(x, x)
+    end
+
+    @testset "powers of tuple products in streams" begin
+        Pt = productmeasure((StdNormal(), StdExponential()^2))
+        m = mcombine(vcat, StdNormal(), Pt^2)
+        x = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]
+        ℓ = logdensityof(StdNormal(), 0.1) + logdensityof(Pt, (0.2, [0.3, 0.4])) + logdensityof(Pt, (0.5, [0.6, 0.7]))
+        @test logdensityof(m, x) ≈ ℓ
+        X = hcat(x, 2 .* x)
+        @test logdensities(m, X) ≈ [logdensityof(m, X[:, j]) for j in 1:2]
+        f = transport_to(StdUniform()^7, m)
+        @test inverse(f)(f(x)) ≈ x
+        @test flatview(inverse(f).(f.(sliced(X, Val(1))))) ≈ X
+        Z, R = batched_transport_to_std_with_rest(StdUniform, Pt, hcat(x, x), (2,))
+        @test size(Z) == (6, 2) && size(R) == (1, 2)
+        Xr, _ = batched_transport_from_std_with_rest(StdUniform, Pt, Z, (2,))
+        @test Xr[1] ≈ [0.1 0.1; 0.4 0.4] && size(Xr[2]) == (2, 2, 2)
+    end
+
+    @testset "batched kernels of measures without a declared rank" begin
+        w = weightedmeasure(0.3, UnknownRankMeasure())
+        @test logdensityof(w, [1.0, 2.0]) ≈ 0.3 - 3
+        @test_throws ArgumentError MeasureBase.batched_logdensity_def(w, randn(2, 2))
+        @test MeasureBase.batched_logdensity_def(weightedmeasure(0.3, StdNormal()^2), [1.0, 2.0]) === 0.3
+    end
+
+    @testset "ragged batches are evaluated variate by variate" begin
+        V = VectorOfVectors([randn(2) for _ in 1:4])
+        @test logdensities(StdNormal()^2, V) ≈ logdensityof.(Ref(StdNormal()^2), V)
+        @test_throws ArgumentError logdensities(StdNormal(), VectorOfVectors([[1.0], [2.0], [3.0], [4.0]]))
+        @test logdensities(StdNormal()^2, nestedview(randn(2, 4))) isa AbstractVector
+    end
+
+    @testset "static streams stay allocation-free" begin
+        m1 = mcombine(vcat, StdNormal(), StdExponential()^static(2))
+        x1 = SVector(0.1, 0.2, 0.3)
+        @test logdensityof(m1, x1) ≈ logdensityof(StdNormal(), 0.1) + logdensityof(StdExponential()^2, [0.2, 0.3])
+        @test @allocated(logdensityof(m1, x1)) == 0
+        m2 = mcombine(vcat, StdNormal()^static(2), StdExponential()^static(3))
+        x2 = SVector(0.1, 0.2, 0.3, 0.4, 0.5)
+        @test @allocated(logdensityof(m2, x2)) == 0
+        g = transport_to(StdUniform()^3, StdNormal()^3)
+        v = randn(3)
+        @test g.(SVector{3}(v))[] ≈ g(v)
+        gs = transport_to(StdUniform(), StdNormal())
+        @test gs.(SVector{3}(v)) isa SVector{3,Float64}
+    end
+
+    @testset "random variates of fused array products" begin
+        P = productmeasure([pushfwd(Base.Fix1(*, s), StdExponential()) for s in 0.1:0.2:0.9])
+        x = rand(GenContext{Float64}(), P)
+        @test x isa Vector{Float64} && length(x) == 5
+        X = batched_rand_impl(GenContext{Float64}(), P, (7,))
+        @test size(X) == (5, 7)
+        Pj = productmeasure(JLArray([weightedmeasure(log(i), StdNormal()) for i in 1:3]))
+        @test length(rand(Pj)) == 3
+    end
+end
diff --git a/test/logdensities.jl b/test/logdensities.jl
index ee4f1a34..c4e57fc1 100644
--- a/test/logdensities.jl
+++ b/test/logdensities.jl
@@ -191,7 +191,7 @@ MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s)
         @test ℓ ≈ vec(sum(stdnormal_ld.(X[1:2, :]), dims = 1))
         @test size(A_rest) == (3, 6)
         ℓ2, A_rest2 = MeasureBase.batched_logdensityof_with_rest(StdNormal(), X, (2,))
-        @test ℓ2 ≈ stdnormal_ld.(X[1:2, :]) && size(A_rest2) == (3, 6)
+        @test MeasureBase._materialize(ℓ2) ≈ stdnormal_ld.(X[1:2, :]) && size(A_rest2) == (3, 6)
         @test_throws ArgumentError logdensities(m, vcat(X, rand(1, 6)))
 
         m3 = mcombine(vcat, StdNormal(), mcombine(vcat, StdExponential()^2, StdLogistic()))
diff --git a/test/runtests.jl b/test/runtests.jl
index 4945ad06..54988f49 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -18,6 +18,7 @@ include("getdof.jl")
 include("shape_contract.jl")
 include("logdensities.jl")
 include("structured_batches.jl")
+include("batched_regressions.jl")
 include("numtype.jl")
 include("transport.jl")
 include("transport_batched.jl")

From a65ce4e878c9ccd6fb3538b12ffdfe67dff10d22 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 13:15:42 +0200
Subject: [PATCH 098/122] Add redesign notes for the major upgrade

Working notes on the batched-first approach, its concepts, extension
points and layout rules, the changes relative to master, device
verification, known limitations and open decisions, to guide reviews and
next steps while the branch evolves. To be removed before the merge.

Created by generative AI.
---
 redesign.md | 173 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 173 insertions(+)
 create mode 100644 redesign.md

diff --git a/redesign.md b/redesign.md
new file mode 100644
index 00000000..333260be
--- /dev/null
+++ b/redesign.md
@@ -0,0 +1,173 @@
+# MeasureBase redesign notes (branch `major-upgrade`)
+
+Working notes on the approach behind this branch, for reviews and for
+guiding the next steps. Kept up to date while the branch evolves, to be
+removed before the merge.
+
+## Goals
+
+- A breaking release that runs on GPUs (CUDA, JLArrays) and under
+  Reactant, with batching built into the foundation.
+- One implementation per measure type for densities, transports and
+  random variates, so scalar and batched paths can't drift apart.
+- Composable, structural solutions: powers, products, combinations,
+  binds and pushforwards implement their behavior once in terms of their
+  components. No shape inference, no per-call DOF sums, no function
+  traits in the core.
+
+## Design philosophy
+
+- **Batched first.** A single variate is a batch with zero batch
+  dimensions. Every kernel handles both, the point API is the batched API
+  at zero batch dimensions. Static arrays keep the scalar path
+  allocation-free.
+- **Ranks, not sizes.** A flat batch is an array
+  `(variate dims..., batch dims...)`. Kernels only need the variate rank
+  of their measure; sizes are optional declarations for validation,
+  stream consumption and static fast paths, never for routing. Unknown
+  sizes are safe.
+- **Standard measures as pivots.** Transport goes through a standard
+  measure type chosen by promoting the measures' preferences. Measure
+  types implement transport to and from their preferred standard measure
+  only.
+- **Branch-free, device-friendly kernels.** Broadcasts, reductions and
+  masks instead of branches; host loops only where documented.
+- **Entry points normalize layouts.** Users pass nested arrays, tuples of
+  batches, struct arrays or flat arrays; the kernels only see flat
+  batches.
+
+## Concepts
+
+**Variate rank.** `mspace_ndims(::Type{M})`, 0 for scalar variates.
+Derived from `mspace_flatsize(::Type{M})` where known, declared by
+array-variate leaves, derived by structural measures. Without a rank the
+batched defaults throw with a message naming the declaration; point
+kernels keep working.
+
+**Flat batches.** Kernels return arrays over the batch dimensions, a
+number for a single variate, possibly lazily. ArraysOfArrays containers
+are fused into their flat storage at the entry points, ragged
+containers are evaluated variate by variate. Tuple and named tuple
+variates batch as tuples of batches; struct arrays and arrays of tuples
+are accepted, their flat storage is the tuple of component storages.
+`rand(Pt^n)` of a tuple product is a struct array.
+
+**Streams.** Variates of `mcombine(vcat, ...)`, binds and tuple products
+inside such streams are flat vectors consumed with the with-rest
+protocol. Point forms return `(result, x_μ, x_rest)` (binds need the
+consumed variate), batched forms take streams `(rows, batch dims...)`
+and a multiplicity `sz::Dims` of variates per stream and return
+`(result, rest)`. Powers pass their size as multiplicity to their base;
+combined measures and tuple products split rows by their fixed stream
+lengths. `fixed_stream_size(::Type{M})` decides whether a batch of
+streams is consumed in fused operations or stream by stream by the
+outermost combinator (binds never fuse). Scalar leaves consume
+`(1, sz..., batch dims...)` and drop the leading dimension. Nested
+element variates in vcat streams are flattened.
+
+**Transport.** `transport_to(ν, μ)` gives a `TransportFunction`; `f(x)`
+transports a variate, `f.(X)` a whole batch. Leaves implement
+`transport_to_std`/`transport_from_std` for their preferred standard
+measure (`preferred_stdmeasure`, `promote_stdmeasure`, `AnyStdMeasure`,
+`NoStdTransport`), array-variate leaves also the `batched_` forms.
+`stdconvert` converts between standard measures in log form,
+`transport_def` may be specialized for direct pairs. Standard streams are
+`(dof, batch dims...)`; the from-side default consumes `fast_dof(μ)`
+entries per variate, composed measures implement the with-rest forms.
+`getdof`/`fast_dof` are declaration-derived, used at construction time
+and for chunking, never inside kernels.
+
+**Random variates.** `rand(ctx::GenContext, μ)` with RNG, precision and
+compute unit (`rand(μ)`, `rand(rng, μ)`, `rand(T, μ)` are wrappers).
+`batched_rand_impl(ctx, μ, sz::Dims)` returns a flat batch, a single
+variate for `sz == ()`; `rand_impl` defaults to it. Defaults draw
+standard variates in bulk on the compute unit and transport them, or
+generate variate by variate without a standard transport.
+
+**Array products.** `productmeasure(::AbstractArray)` stores isbits
+parameterized marginals as `StructArrays` (nested parameter structs
+unwrapped; numbers, arrays, tuples, strings, symbols and function objects
+stay opaque columns), in one place: `_marginal_storage`. Fused kernels
+broadcast over the leaf columns and rebuild marginals via
+`ConstructionBase.constructorof`, which works on CUDA and under Reactant.
+Fusion needs concrete scalar-variate marginals (one DOF for transport);
+other array products loop over host-resident marginals. Measures holding
+arrays have `Adapt` rules.
+
+**Pushforwards.** `pushfwd(f, μ)` learns its output size once at
+construction from a test value when the origin has a size. Batched
+application needs no traits: elementwise for `Base.BroadcastFunction`
+(fused density kernels), column batches for AffineMaps types (weak
+dependency, `MeasureBaseAffineMapsExt`), a host loop otherwise.
+
+## Extension points
+
+| Aspect | Scalar-variate leaf | Array-variate leaf | Composed measure |
+|---|---|---|---|
+| Density | `logdensity_def` | `mspace_ndims`, `batched_logdensityof_impl` (+`_def`) | kernels via components, with-rest forms |
+| Transport | `preferred_stdmeasure`, `transport_to_std`, `transport_from_std` | + `batched_transport_to_std`, `batched_transport_from_std` | with-rest forms, `fixed_stream_size` |
+| Random | `batched_rand_impl` (default via transport) | `batched_rand_impl` | derived |
+| Declarations | none | `mspace_ndims`, optionally `mspace_flatsize`, `getdof` | derived |
+
+## Layouts per combinator
+
+- Powers: `(base dims..., power dims..., batch dims...)`, innermost base
+  first; standard variates are the flat vector of the base's. Results
+  are nested views over the flat storage. Powers of tuple products treat
+  numeric flat variates as streams and accept tuples of batches.
+- Array products: `(marginal dims..., product dims..., batch dims...)`.
+- Tuple products: tuples of batches; marginal by marginal in streams.
+- Combined `vcat`: streams; `merge`: merged named tuples. Binds: single
+  streams, value-dependent sizes.
+- Weighted, restricted, density measures, Half: forward plus weights or
+  masks. Superpositions and spike mixtures: one batch per component,
+  masks aligned with the variate dimensions. Dirac: constant batches.
+- Distributions extension: univariate via `StdLogistic` (log-cdf and
+  quantile), location-scale families via their affine map, `MvNormal` via
+  Cholesky factors, array-variate batches via `logpdf(d, X)`.
+
+## Changes relative to `master`
+
+Twenty commits since `c773afe`: variate size contract and
+`preferred_stdmeasure`, densities over flat storage, branch-free
+evaluation, Reactant smoke tests, structural batched kernels, transport
+rebuilt on standard measures, batched transport and the broadcast hook,
+rand via `GenContext`, then the batched-first redesign (density core,
+struct array products, transport, random variates, structured batches)
+and the review fixes.
+
+Removed: `transport_origin`/`to_origin`/`from_origin` and the origin
+machinery, `NoTransport`, `transport_to_mvstd`, per-measure
+`Base.rand(rng, T, μ)` methods, the `_trafo_cdf`/`_trafo_quantile` hooks.
+
+Behavior changes for NEWS: univariate Distributions pivot on
+`StdLogistic`; nested powers and array products return ArraysOfArrays
+views, powers of tuple products struct arrays; `Half` transports via
+`StdUniform` (tails limited); `mcombine(vcat, ::Product, ::Product)`
+merges only concrete homogeneous marginals; `rand(rng, Int, μ)`
+unsupported; isbits marginal arrays become struct arrays (`rand` of such
+products gives plain arrays); vcat-combined and bind variates are flat;
+`transport_to(StdUniform, m)` with binds inside errors (use
+`StdUniform()^n`); batched kernels of rank-less measures throw.
+
+## Verification
+
+Full suite (Aqua, extensions, doctests) on CPU with JLArrays cases;
+`test/reactant` (opt-in, backend via `MEASUREBASE_REACTANT_BACKEND`) and
+`test/cuda` (opt-in) run locally on the GB10, both green at HEAD except
+one expected-broken CUDA case (AffineMaps Jacobian on device).
+
+## Known gaps and open decisions
+
+- Upstream: AffineMaps lacks `Adapt` rules and device/traced Jacobians;
+  Distributions isn't device-aware; ChangesOfVariables has no rules for
+  `Base.Fix1`/`Fix2` arithmetic; HeterogeneousComputing has no Reactant
+  compute unit; JLArrays has no RNG; Reactant rejects traced
+  `VectorOfArrays` and empty batches.
+- Decisions pending: one convention for out-of-support inputs (NaN mask
+  vs. DomainError vs. AssertionError), `Half` tails via log-ccdf, device
+  random variate infrastructure and `rand!`, Tier-1 static variates,
+  the `smart-constructors.jl` review (location-scale arrays as affine
+  pushforwards of powers), `_static_ndims` type-first vs. instance-first.
+- Polish before merge: docs pass, NEWS, history curation, version bump,
+  remove this file.

From 92b08171f6c893e0e341f23103f5547bfe04c4c9 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 14:26:37 +0200
Subject: [PATCH 099/122] Make Reactant an on-demand test dependency

Reactant only supports 64-bit Linux and macOS, so runtests.jl adds it
on the fly where supported (as MGVI does) and includes the Reactant
tests from test/test_reactant.jl instead of a separate test project.

Created by generative AI.
---
 test/reactant/Project.toml                    | 12 -----
 test/runtests.jl                              |  9 ++++
 .../runtests.jl => test_reactant.jl}          | 44 ++++++++++++++++---
 3 files changed, 46 insertions(+), 19 deletions(-)
 delete mode 100644 test/reactant/Project.toml
 rename test/{reactant/runtests.jl => test_reactant.jl} (69%)

diff --git a/test/reactant/Project.toml b/test/reactant/Project.toml
deleted file mode 100644
index 775cd85e..00000000
--- a/test/reactant/Project.toml
+++ /dev/null
@@ -1,12 +0,0 @@
-[deps]
-ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
-Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
-MeasureBase = "fa1605e6-acd5-459c-a1e6-7e635759db14"
-Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[sources]
-MeasureBase = {path = "../.."}
-
-[compat]
-Reactant = "0.2"
diff --git a/test/runtests.jl b/test/runtests.jl
index 54988f49..55b9d679 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -44,4 +44,13 @@ include("rand_batched.jl")
 
 include("distributions/test_distributions.jl")
 
+# Reactant only supports 64-bit Linux and macOS, and some of its
+# dependencies break already during precompilation on other platforms,
+# so it can't be a static test dependency:
+if Sys.WORD_SIZE == 64 && (Sys.islinux() || Sys.isapple()) && isempty(VERSION.prerelease)
+    import Pkg
+    Base.identify_package("Reactant") === nothing && Pkg.add("Reactant")
+    include("test_reactant.jl")
+end
+
 include("test_docs.jl")
diff --git a/test/reactant/runtests.jl b/test/test_reactant.jl
similarity index 69%
rename from test/reactant/runtests.jl
rename to test/test_reactant.jl
index 0fee2da1..4df1f7db 100644
--- a/test/reactant/runtests.jl
+++ b/test/test_reactant.jl
@@ -1,19 +1,21 @@
 # This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
 
-# Reactant smoke tests, not part of the default test suite. Run with
-# `julia --project=test/reactant test/reactant/runtests.jl` after
-# instantiating that project, or include this file in an environment that
-# provides Reactant. The backend defaults to the CPU, set the environment
-# variable `MEASUREBASE_REACTANT_BACKEND` (e.g. to "gpu") to change it.
+# Reactant tests. Reactant isn't a static test dependency (it only
+# supports 64-bit Linux and macOS), runtests.jl adds it on the fly where
+# supported. The backend defaults to the CPU, set the environment variable
+# `MEASUREBASE_REACTANT_BACKEND` (e.g. to "gpu") to change it; the file can
+# also be run standalone in an environment that provides Reactant.
 
 using Test
 using Reactant
 using MeasureBase
-using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic, Lebesgue, Dirac
+using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic, Lebesgue, Dirac, asmeasure
 using MeasureBase: logdensities, logdensity_rel, weightedmeasure, superpose, restrict, mintegrate_exp
 using MeasureBase: mcombine
 using ArraysOfArrays: VectorOfSimilarVectors, sliced, flatview
-using Distributions: Normal, Exponential, Uniform, Beta
+using Distributions: Normal, Uniform, Exponential, Logistic, Cauchy, Laplace, LogNormal, Weibull, Gamma, Beta
+using Distributions: Poisson, MvNormal, Dirichlet
+using MeasureBase.InverseFunctions: inverse
 
 Reactant.set_default_backend(get(ENV, "MEASUREBASE_REACTANT_BACKEND", "cpu"))
 
@@ -107,4 +109,32 @@ _plain(x::Number) = Float64(x)
         test_traced(x -> transport_to(Normal(2, 3), StdNormal()).(x), x)
         test_traced(x -> transport_to(StdNormal(), Exponential(2.0)).(x), rand(10))
     end
+
+    # Distribution parameters stay constants, Distributions' parameter
+    # structs can't hold traced arrays:
+    @testset "wrapped distributions" begin
+        for d in (Normal(0.3, 1.7), Uniform(-1.0, 2.5), Exponential(0.7), Logistic(0.2, 1.3), Cauchy(0.1, 0.8), Laplace(-0.4, 1.1), LogNormal(0.2, 0.6), Weibull(1.4, 0.9), Gamma(2.3, 1.2), Beta(2.5, 3.5))
+            m = asmeasure(d)
+            xd = rand(d, 10)
+            test_traced(x -> logdensities(m, x), xd)
+            f = transport_to(StdNormal(), m)
+            if d isa Union{Gamma,Beta}
+                # SpecialFunctions' incomplete gamma and beta functions have no Reactant methods:
+                @test_broken @jit((x -> copy(f.(x)))(Reactant.to_rarray(xd))) isa AbstractArray
+            else
+                test_traced(x -> f.(x), xd)
+                test_traced(z -> inverse(f).(z), randn(10))
+            end
+        end
+        test_traced(x -> logdensities(asmeasure(Poisson(2.7)), x), Float64.(rand(Poisson(2.7), 10)))
+        mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3])
+        mm = asmeasure(mvn)
+        Xm = rand(mvn, 10)
+        test_traced(X -> logdensities(mm, X), Xm)
+        test_traced(X -> flatview(transport_to(StdNormal()^2, mm).(sliced(X, Val(1)))), Xm)
+        test_traced(Z -> flatview(transport_to(mm, StdNormal()^2).(sliced(Z, Val(1)))), randn(2, 10))
+        dir = Dirichlet([2.0, 3.0, 4.0, 1.5])
+        md = asmeasure(dir)
+        test_traced(X -> logdensities(md, X), rand(dir, 10))
+    end
 end

From 3a1d310bd7313ce1013e7236f34aebe71199a2bf Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 14:26:37 +0200
Subject: [PATCH 100/122] Fix CI failures on Julia 1.10 and 1.13 and in the
 docs build

Leading singleton dimensions of summed batches are dropped by reshaping,
since dropdims on static matrices doesn't infer on Julia 1.13. The
common variate rank of superposition components is folded pairwise and
density measures evaluate array variates directly, which Julia 1.10
needs to infer these kernels. The allocation tests call the kernels
directly, local closures allocate inside test sets on Julia 1.10. The
docs include the MeasureOperators docstrings and all cross-references
resolve.

Created by generative AI.
---
 docs/make.jl                        |  2 +-
 docs/src/api.md                     |  2 +-
 src/combinators/implicitlymapped.jl |  6 +++---
 src/combinators/superpose.jl        |  9 ++++++---
 src/density-batched.jl              |  9 ++++++++-
 src/density.jl                      |  4 ++++
 src/measure_operators.jl            |  2 +-
 src/mspace.jl                       |  4 ++--
 src/primitives/dirac.jl             |  2 +-
 src/standard/stdexponential.jl      |  6 ++++++
 src/standard/stdlogistic.jl         |  6 ++++++
 src/standard/stdnormal.jl           |  6 ++++++
 src/standard/stduniform.jl          |  6 ++++++
 test/logdensities.jl                | 20 ++++++++------------
 14 files changed, 59 insertions(+), 25 deletions(-)

diff --git a/docs/make.jl b/docs/make.jl
index 3407cf93..96f147eb 100644
--- a/docs/make.jl
+++ b/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/",
diff --git a/docs/src/api.md b/docs/src/api.md
index 83c27f70..fc4a8a33 100644
--- a/docs/src/api.md
+++ b/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]
 ```
diff --git a/src/combinators/implicitlymapped.jl b/src/combinators/implicitlymapped.jl
index 3966b10a..36a72a93 100644
--- a/src/combinators/implicitlymapped.jl
+++ b/src/combinators/implicitlymapped.jl
@@ -84,7 +84,7 @@ export ImplicitlyMapped
 Get the original object (a measure or transition/Markov kernel) that was
 implicitly mapped.
 
-See [ImplicitlyMapped](@ref) for detailed semantics.
+See [`ImplicitlyMapped`](@ref) for detailed semantics.
 
 # Implementation
 
@@ -100,7 +100,7 @@ export implicit_origin
 Get an explicit map/function based on an implicitly mapped object and an
 observation.
 
-See [ImplicitlyMapped](@ref) for detailed semantics.
+See [`ImplicitlyMapped`](@ref) for detailed semantics.
 
 # Implementation
 
@@ -206,7 +206,7 @@ Constructors:
 * `Marginalized(mu)`
 * `Marginalized(f_kernel)`
 
-See [ImplicitlyMapped](@ref) for detailed semantics.
+See [`ImplicitlyMapped`](@ref) for detailed semantics.
 
 Example:
 
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index d0370eed..9042d283 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -179,10 +179,13 @@ end
     args = [:(mspace_ndims($T)) for T in C.parameters]
     :(_common_ndims(($(args...),), MU))
 end
-@inline function _common_ndims(ns::Tuple{Integer,Vararg{Integer}}, ::Type{MU}) where {MU}
-    all(==(first(ns)), ns) ? first(ns) : NoMSpaceElementSize{MU}()
-end
+# Pairwise comparisons fold to a constant rank where `all` doesn't (Julia 1.10):
+@inline _common_ndims(ns::Tuple{Integer,Vararg{Integer}}, ::Type{MU}) where {MU} = _common_ndims_of(first(ns), Base.tail(ns), MU)
 @inline _common_ndims(::Tuple, ::Type{MU}) where {MU} = NoMSpaceElementSize{MU}()
+@inline _common_ndims_of(n::Integer, ::Tuple{}, ::Type) = n
+@inline function _common_ndims_of(n::Integer, ns::Tuple{Integer,Vararg{Integer}}, ::Type{MU}) where {MU}
+    n == first(ns) ? _common_ndims_of(n, Base.tail(ns), MU) : NoMSpaceElementSize{MU}()
+end
 @inline mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:AbstractArray} = _scalar_or_unknown(mspace_flatsize(eltype(C)))
 @inline mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:Tuple} = _common_scalar_flatsize(C)
 @generated function _common_scalar_flatsize(::Type{C}) where {C<:Tuple}
diff --git a/src/density-batched.jl b/src/density-batched.jl
index 91d9f6df..e5748699 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -235,7 +235,14 @@ end
 @inline _sum_leading_dims_impl(A::AbstractArray, ::StaticInteger{0}, ::StaticInteger{0}) = A
 @inline _sum_leading_dims_impl(A::AbstractArray, ::StaticInteger{N}, ::StaticInteger{N}) where {N} = sum(A)
 @inline function _sum_leading_dims_impl(A::AbstractArray, ::StaticInteger{N}, ::StaticInteger) where {N}
-    dropdims(_sum_dims_seq(A, static(N)); dims = ntuple(identity, Val(N)))
+    _drop_leading_dims(_sum_dims_seq(A, static(N)), static(N))
+end
+
+# Drops the leading `N` (singleton) dimensions by reshaping, which keeps
+# static arrays static and infers where `dropdims` doesn't:
+@inline function _drop_leading_dims(A::AbstractArray, ::StaticInteger{N}) where {N}
+    dims = _batch_dims(A)
+    _reshape_batch(A, ntuple(i -> dims[N + i], Val(length(dims) - N)))
 end
 @inline _sum_dims_seq(A::AbstractArray, ::StaticInteger{0}) = A
 @inline function _sum_dims_seq(A::AbstractArray, ::StaticInteger{N}) where {N}
diff --git a/src/density.jl b/src/density.jl
index 361b0851..f6dabe46 100644
--- a/src/density.jl
+++ b/src/density.jl
@@ -230,6 +230,10 @@ logdensity_def(μ::DensityMeasure, x) = logdensityof(μ.f, x)
 
 density_def(μ::DensityMeasure, x) = densityof(μ.f, x)
 
+# Density measures evaluate the base measure and the integrand directly,
+# the base measure validates array variates:
+@inline _point_ld(f::F, μ::DensityMeasure, x::AbstractArray{<:Number}) where {F} = f(μ, x)
+
 function logdensityof_impl(μ::DensityMeasure, x::Any)
     integrand, μ_base = μ.f, μ.base
     base_logval = dynamic(logdensityof(μ_base, x))
diff --git a/src/measure_operators.jl b/src/measure_operators.jl
index 41606367..42fbb837 100644
--- a/src/measure_operators.jl
+++ b/src/measure_operators.jl
@@ -52,7 +52,7 @@ A common mathematical notation for pullback in measure theory is
 ``f \circ μ``, but as `∘` is used for function composition in Julia and as
 `f` semantically acts point-wise on sets, we use `⊙`.
 
-Also see [f ⋄ μ](@ref), the pushforward operator.
+Also see `f ⋄ μ`, the pushforward operator.
 """
 ⊙(ν::AbstractMeasure, f) = pullbck(f, ν)
 export ⊙
diff --git a/src/mspace.jl b/src/mspace.jl
index c1520863..4ddb28ef 100644
--- a/src/mspace.jl
+++ b/src/mspace.jl
@@ -15,7 +15,7 @@ For a measure `μ` over an array-valued measurable space, return the size of
 the arrays that are the elements of the space, `()` for scalar variates.
 
 The size is static where it is known statically. Returns
-[`NoMSpaceElementSize{typeof(μ)}()`](@ref) if the elements of the space
+[`NoMSpaceElementSize{typeof(μ)}()`](@ref NoMSpaceElementSize) if the elements of the space
 are not arrays of one common size, e.g. for structured variates or variates
 whose size depends on the value, or if the size can not be determined
 efficiently.
@@ -36,7 +36,7 @@ variates.
 
 Variates of powers of measures with array-valued variates are nested
 arrays, their flat storage has the size of the inner arrays followed by
-the size of the power. Returns [`NoMSpaceElementSize{typeof(μ)}()`](@ref)
+the size of the power. Returns [`NoMSpaceElementSize{typeof(μ)}()`](@ref NoMSpaceElementSize)
 if the variates of `μ` have no flat storage of a common size.
 
 See also [`mspace_elsize`](@ref).
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 5a92d639..08a27a83 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -91,7 +91,7 @@ end
 
 @inline _all_leading_dims(A::AbstractArray{Bool,N}, ::StaticInteger{N}) where {N} = all(A)
 @inline function _all_leading_dims(A::AbstractArray{Bool}, ::StaticInteger{N}) where {N}
-    dropdims(all(A; dims = ntuple(identity, Val(N))); dims = ntuple(identity, Val(N)))
+    _drop_leading_dims(all(A; dims = ntuple(identity, Val(N))), static(N))
 end
 
 Adapt.adapt_structure(to, μ::Dirac) = Dirac(Adapt.adapt(to, μ.x))
diff --git a/src/standard/stdexponential.jl b/src/standard/stdexponential.jl
index 82abc200..39043e4a 100644
--- a/src/standard/stdexponential.jl
+++ b/src/standard/stdexponential.jl
@@ -1,3 +1,9 @@
+"""
+    StdExponential <: StdMeasure
+
+The standard exponential measure, the exponential distribution with unit
+scale as a measure.
+"""
 struct StdExponential <: StdMeasure end
 
 export StdExponential
diff --git a/src/standard/stdlogistic.jl b/src/standard/stdlogistic.jl
index fa6415a0..44b8a0b8 100644
--- a/src/standard/stdlogistic.jl
+++ b/src/standard/stdlogistic.jl
@@ -1,3 +1,9 @@
+"""
+    StdLogistic <: StdMeasure
+
+The standard logistic measure, the logistic distribution with zero
+location and unit scale as a measure.
+"""
 struct StdLogistic <: StdMeasure end
 
 export StdLogistic
diff --git a/src/standard/stdnormal.jl b/src/standard/stdnormal.jl
index faa043dc..989e645e 100644
--- a/src/standard/stdnormal.jl
+++ b/src/standard/stdnormal.jl
@@ -1,6 +1,12 @@
 using SpecialFunctions: erfc, erfcinv, logerfc
 using IrrationalConstants: invsqrt2, log2π, logtwo
 
+"""
+    StdNormal <: StdMeasure
+
+The standard normal measure, the normal distribution with zero mean and
+unit variance as a measure.
+"""
 struct StdNormal <: StdMeasure end
 
 export StdNormal
diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl
index 7c66caec..b5ec1afa 100644
--- a/src/standard/stduniform.jl
+++ b/src/standard/stduniform.jl
@@ -1,3 +1,9 @@
+"""
+    StdUniform <: StdMeasure
+
+The standard uniform measure on the unit interval, the uniform
+distribution on `[0, 1]` as a measure.
+"""
 struct StdUniform <: StdMeasure end
 
 export StdUniform
diff --git a/test/logdensities.jl b/test/logdensities.jl
index c4e57fc1..cfa32e36 100644
--- a/test/logdensities.jl
+++ b/test/logdensities.jl
@@ -124,13 +124,11 @@ MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s)
     @testset "static variates" begin
         m3 = StdNormal()^static(3)
         xs = @SVector randn(3)
-        f(x) = logdensityof(m3, x)
-        @test @inferred(f(xs)) ≈ sum(stdnormal_ld, xs)
-        @test @allocated(f(xs)) == 0
-        g(x) = logdensityof(StdNormal()^3, x)
+        @test @inferred(logdensityof(m3, xs)) ≈ sum(stdnormal_ld, xs)
+        @test @allocated(logdensityof(m3, xs)) == 0
         xd = randn(3)
-        @test @inferred(g(xd)) ≈ sum(stdnormal_ld, xd)
-        @test @allocated(g(xd)) == 0
+        @test @inferred(logdensityof(StdNormal()^3, xd)) ≈ sum(stdnormal_ld, xd)
+        @test @allocated(logdensityof(StdNormal()^3, xd)) == 0
         Xs = @SMatrix randn(3, 4)
         @test @inferred(logdensities(m3, Xs)) ≈ vec(sum(stdnormal_ld.(Xs), dims = 1))
         @test logdensities(m3, Xs) isa SVector{4}
@@ -149,18 +147,16 @@ MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s)
         X = randn(3, 10)
         @test @inferred(logdensities(w, X)) ≈ [logdensityof(w, x) for x in eachcol(X)]
         xw = randn(3)
-        fw(x) = logdensityof(w, x)
-        @test @inferred(fw(xw)) ≈ log(0.3) + sum(stdnormal_ld, xw)
-        @test @allocated(fw(xw)) == 0
+        @test @inferred(logdensityof(w, xw)) ≈ log(0.3) + sum(stdnormal_ld, xw)
+        @test @allocated(logdensityof(w, xw)) == 0
 
         ms = [weightedmeasure(log(i), StdNormal()) for i in 1:4]
         prod4 = productmeasure(ms)
         @test @inferred(MeasureBase.mspace_flatsize(prod4)) == (4,)
         @test @inferred(MeasureBase.mspace_elsize(prod4)) == (4,)
         xp = randn(4)
-        fp(x) = logdensityof(prod4, x)
-        @test @inferred(fp(xp)) ≈ sum(log(i) + stdnormal_ld(xp[i]) for i in 1:4)
-        @test @allocated(fp(xp)) == 0
+        @test @inferred(logdensityof(prod4, xp)) ≈ sum(log(i) + stdnormal_ld(xp[i]) for i in 1:4)
+        @test @allocated(logdensityof(prod4, xp)) == 0
         Xp = randn(4, 7)
         @test @inferred(logdensities(prod4, Xp)) ≈ [logdensityof(prod4, x) for x in eachcol(Xp)]
         @test logdensities(prod4, sliced(Xp, 1)) ≈ logdensities(prod4, Xp)

From 895ff0264985637224c3d5c395c277bcd52a6f7a Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 14:26:37 +0200
Subject: [PATCH 101/122] Device-friendly kernels for wrapped Distributions

The main univariate families (Normal, Uniform, Exponential, Logistic,
Cauchy, Laplace, LogNormal, Weibull, Gamma, Beta, Poisson, Bernoulli)
get branch-free density formulas, support masks and transports to the
standard measure matching their shape, so batches evaluate on GPUs and
under Reactant. MvNormal works on column batches through its Cholesky
factor (kept lazy, Cholesky.L scalar-indexes on GPUs) and Dirichlet
through stick-breaking Beta transports with cumulative sums and
products along the variate dimension, both with Adapt rules. Incomplete
gamma and beta functions go through hooks whose ForwardDiff and
ChainRules derivatives come from the densities. Draws use the
Distributions samplers on the CPU and the standard transports on other
compute units. Batched densities of univariate wrappers broadcast their
point kernels, only array-variate wrappers evaluate via logpdf.

Beta and Dirichlet transports stay on the CPU (SpecialFunctions'
incomplete beta functions don't compile for GPUs or Reactant), MvNormal
and Dirichlet parameters stay constants under Reactant.

Created by generative AI.
---
 ...asureBaseDistributionsChainRulesCoreExt.jl |  25 ++++
 .../MeasureBaseDistributionsExt.jl            |  12 +-
 ext/MeasureBaseDistributionsExt/dirichlet.jl  |  51 -------
 .../distribution_measure.jl                   |  24 +++-
 ext/MeasureBaseDistributionsExt/families.jl   | 133 ++++++++++++++++++
 .../multivariate.jl                           | 112 +++++++++++++++
 ext/MeasureBaseDistributionsExt/standardmv.jl |  25 +---
 ext/MeasureBaseDistributionsForwardDiffExt.jl |  44 ++++++
 redesign.md                                   |  45 ++++--
 src/utils.jl                                  |  21 +++
 test/cuda/runtests.jl                         |  37 ++++-
 test/distributions/test_device_kernels.jl     | 114 +++++++++++++++
 test/distributions/test_distributions.jl      |   1 +
 test/distributions/test_shape_contract.jl     |   6 +-
 14 files changed, 548 insertions(+), 102 deletions(-)
 create mode 100644 ext/MeasureBaseDistributionsExt/families.jl
 create mode 100644 ext/MeasureBaseDistributionsExt/multivariate.jl
 create mode 100644 test/distributions/test_device_kernels.jl

diff --git a/ext/MeasureBaseDistributionsChainRulesCoreExt.jl b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl
index 63cc3e93..02fa145d 100644
--- a/ext/MeasureBaseDistributionsChainRulesCoreExt.jl
+++ b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl
@@ -11,6 +11,31 @@ 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
diff --git a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
index 83a20c65..7795e25a 100644
--- a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
+++ b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
@@ -2,7 +2,8 @@
 
 module MeasureBaseDistributionsExt
 
-using LinearAlgebra: Diagonal, diag, dot, cholesky
+using LinearAlgebra: Diagonal, Cholesky, LowerTriangular, UpperTriangular, diag, dot, cholesky
+import Adapt
 
 import Random
 using Random: AbstractRNG, rand!
@@ -25,11 +26,13 @@ import MeasureBase:
     _dist_params_numtype, _trafo_logcdf_impl, _trafo_logccdf_impl,
     _trafo_quantile_impl, _trafo_cquantile_impl, _dist_quantile, _dist_cquantile
 using MeasureBase: _pushfront, _pushback, _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log
+using MeasureBase: _gamma_cdf, _gamma_quantile, _beta_cdf, _beta_quantile, _gamma_logpdf, _beta_logpdf, _dualtag
 
 import Distributions
 using Distributions: Distribution, VariateForm, ValueSupport, ContinuousDistribution
 using Distributions: Univariate, Multivariate, ArrayLikeVariate, Continuous, Discrete
 using Distributions: Uniform, Exponential, Logistic, Normal
+using Distributions: Cauchy, Laplace, LogNormal, Weibull, Gamma, Poisson, Bernoulli
 using Distributions: MvNormal, AbstractMvNormal, Beta, Dirichlet
 using Distributions: ReshapedDistribution, AbstractMixtureModel
 
@@ -39,9 +42,10 @@ import StatsFuns
 import PDMats
 
 using IrrationalConstants: log2π, invsqrt2π
-using LogExpFunctions: logistic
+using LogExpFunctions: logistic, log1pexp
+using SpecialFunctions: loggamma, logbeta, gamma_inc, gamma_inc_inv, beta_inc, beta_inc_inv
 
-using HeterogeneousComputing: real_numtype, GenContext, get_rng, get_precision
+using HeterogeneousComputing: real_numtype, GenContext, get_rng, get_precision, get_compute_unit, CPUnit, AbstractComputeUnit
 
 using Static: True, False, StaticInt, static, dynamic
 using StaticThings: asnonstatic
@@ -58,11 +62,13 @@ include("standard_normal.jl")
 include("distribution_measure.jl")
 include("dist_vartransform.jl")
 include("univariate.jl")
+include("families.jl")
 include("standardmv.jl")
 include("product.jl")
 include("reshaped.jl")
 include("mixture.jl")
 include("dirichlet.jl")
+include("multivariate.jl")
 include("dirac.jl")
 
 end # module MeasureBaseDistributionsExt
diff --git a/ext/MeasureBaseDistributionsExt/dirichlet.jl b/ext/MeasureBaseDistributionsExt/dirichlet.jl
index 90942710..b156d4a7 100644
--- a/ext/MeasureBaseDistributionsExt/dirichlet.jl
+++ b/ext/MeasureBaseDistributionsExt/dirichlet.jl
@@ -6,54 +6,3 @@ MeasureBase.getdof(d::Dirichlet) = length(d) - 1
 MeasureBase.getdof(m::DirichletMeasure) = getdof(m.obj)
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Dirichlet}) = StdUniform
-
-
-
-function _dirichlet_beta_trafo(α::Real, β::Real, x::Real)
-    R = float(promote_type(typeof(α), typeof(β), typeof(x)))
-    convert(R, transport_def(Beta(α, β), StdUniform(), x))::R
-end
-
-_a_times_one_minus_b(a::Real, b::Real) = a * (1 - b)
-
-function MeasureBase.transport_from_std(::Type{StdUniform}, ν::Dirichlet, x)
-    # See M. J. Betancourt, "Cruising The Simplex: Hamiltonian Monte Carlo and the Dirichlet Distribution",
-    # https://arxiv.org/abs/1010.3436
-
-    @_adignore @argcheck length(ν) == length(x) + 1
-
-    αs = _dropfront(_rev_cumsum(ν.alpha))
-    βs = _dropback(ν.alpha)
-    beta_v = _fwddiff(_dirichlet_beta_trafo).(αs, βs, x)
-    beta_v_cp = _exp_cumsum_log(_pushfront(beta_v, 1))
-    beta_v_ext = _pushback(beta_v, 0)
-    _fwddiff(_a_times_one_minus_b).(beta_v_cp, beta_v_ext)
-end
-
-
-function _inv_dirichlet_beta_trafo(α::Real, β::Real, beta_v::Real)
-    R = float(promote_type(typeof(α), typeof(β), typeof(beta_v)))
-    convert(R, transport_def(StdUniform(), Beta(α, β), beta_v))::R
-end
-
-# ToDo: Find efficient pullback for this:
-function _dirichlet_variate_to_beta_v(y::AbstractVector{<:Real})
-    beta_v = similar(y, length(eachindex(y)) - 1)
-    @assert firstindex(beta_v) == firstindex(y)
-    @assert lastindex(beta_v) == lastindex(y) - 1
-    T = eltype(y)
-    sum_log_beta_v::T = 0
-    @inbounds for i in eachindex(beta_v)
-        beta_v[i] = 1 - y[i] / exp(sum_log_beta_v)
-        sum_log_beta_v += log(beta_v[i])
-    end
-    return beta_v
-end
-
-function MeasureBase.transport_to_std(::Type{StdUniform}, ν::Dirichlet, y)
-    @_adignore @argcheck length(ν) == length(y)
-    αs = _dropfront(_rev_cumsum(ν.alpha))
-    βs = _dropback(ν.alpha)
-    beta_v = _dirichlet_variate_to_beta_v(y)
-    _fwddiff(_inv_dirichlet_beta_trafo).(αs, βs, beta_v)
-end
diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index b480e11e..f0c152cb 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -15,11 +15,19 @@ const DistributionMeasure{F<:VariateForm,S<:ValueSupport,D<:Distribution{F,S}} =
 @inline Base.convert(::Type{Distribution{F,S}}, m::DistributionMeasure{F,S}) where {F<:VariateForm,S<:ValueSupport} = Distribution(m)
 
 
-MeasureBase.rand_impl(ctx::GenContext, m::DistributionMeasure) =
-    convert_realtype(get_precision(ctx), rand(get_rng(ctx), m.obj))
+# Distributions' samplers run on the CPU, variates on other compute units
+# are generated from standard variates via the transports:
+MeasureBase.rand_impl(ctx::GenContext, m::DistributionMeasure) = _dist_rand(ctx, m, get_compute_unit(ctx))
+MeasureBase.batched_rand_impl(ctx::GenContext, m::DistributionMeasure, sz::Dims) = _dist_batched_rand(ctx, m, sz, get_compute_unit(ctx))
 
-MeasureBase.batched_rand_impl(ctx::GenContext, m::DistributionMeasure, sz::Dims) =
+_dist_rand(ctx::GenContext, m::DistributionMeasure, ::CPUnit) =
+    convert_realtype(get_precision(ctx), rand(get_rng(ctx), m.obj))
+_dist_rand(ctx::GenContext, m::DistributionMeasure, ::AbstractComputeUnit) =
+    MeasureBase._rand_default(ctx, m, (), MeasureBase._NoRandImpl())
+_dist_batched_rand(ctx::GenContext, m::DistributionMeasure, sz::Dims, ::CPUnit) =
     _flat_powrand(get_rng(ctx), get_precision(ctx), m.obj, sz)
+_dist_batched_rand(ctx::GenContext, m::DistributionMeasure, sz::Dims, ::AbstractComputeUnit) =
+    MeasureBase._rand_default(ctx, m, sz, MeasureBase._NoRandImpl())
 
 # A single variate for zero batch dimensions, flat batches otherwise:
 _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution, ::Tuple{}) where {T<:Real} = convert_realtype(T, rand(rng, d))
@@ -49,11 +57,15 @@ end
 @inline MeasureBase.logdensity_def(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x)
 
 # Distributions evaluate flat batches of array variates (the trailing
-# dimensions are batch dimensions) directly:
-for bhead in (:batched_logdensityof_impl, :batched_logdensity_def)
-    @eval function MeasureBase.$bhead(m::DistributionMeasure{<:ArrayLikeVariate{N}}, X::AbstractArray{<:Real}) where {N}
+# dimensions are batch dimensions) directly, univariate wrappers broadcast
+# their point kernels:
+for (bhead, phead) in ((:batched_logdensityof_impl, :logdensityof_impl), (:batched_logdensity_def, :logdensity_def))
+    @eval function MeasureBase.$bhead(m::DistributionMeasure{<:ArrayLikeVariate{N}}, X::AbstractArray) where {N}
         Distributions.logpdf(m.obj, X)
     end
+    @eval function MeasureBase.$bhead(m::DistributionMeasure{<:ArrayLikeVariate{0}}, X::AbstractArray)
+        MeasureBase._scalar_kernel_broadcast(MeasureBase.$phead, m, X)
+    end
 end
 @inline MeasureBase.unsafe_logdensityof(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x)
 @inline MeasureBase.insupport(m::DistributionMeasure, x) = Distributions.insupport(m.obj, x)
diff --git a/ext/MeasureBaseDistributionsExt/families.jl b/ext/MeasureBaseDistributionsExt/families.jl
new file mode 100644
index 00000000..a0f757e0
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/families.jl
@@ -0,0 +1,133 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+# Densities and standard transports of the main distribution families,
+# implemented as plain arithmetic on the parameters without branches or
+# foreign calls, so that the batched kernels of the wrapped measures run on
+# devices and in traced code. Distributions' own implementations remain in
+# use for other families.
+
+const _Families = Union{Normal,Uniform,Exponential,Logistic,Cauchy,Laplace,LogNormal,Weibull,Gamma,Beta,Poisson,Bernoulli}
+
+# Densities relative to the base measures (Lebesgue resp. counting
+# measure), support checks are separate masks. The density formulas must
+# not throw outside of the support, where their results are masked:
+@inline MeasureBase.logdensity_def(m::AsMeasure{<:_Families}, x) = _family_logd(m.obj, x)
+@inline MeasureBase.unsafe_logdensityof(m::AsMeasure{<:_Families}, x) = _family_logd(m.obj, x)
+@inline MeasureBase.insupport(m::AsMeasure{<:_Families}, x) = _family_insupport(m.obj, x)
+
+# `c * log(y)`, zero for `c == 0` also where `y == 0`:
+@inline _clog(c, y) = ifelse(iszero(c), zero(c * log(one(y))), c * log(y))
+
+@inline function _family_logd(d::Normal, x)
+    z = (x - d.μ) / d.σ
+    -z * z / 2 - log(d.σ) - log2π / 2
+end
+@inline _family_insupport(::Normal, x) = true
+
+@inline _family_logd(d::Uniform, x) = -log(d.b - d.a) + zero(x)
+@inline _family_insupport(d::Uniform, x) = (d.a <= x) & (x <= d.b)
+
+@inline _family_logd(d::Exponential, x) = -x / d.θ - log(d.θ)
+@inline _family_insupport(::Exponential, x) = x >= 0
+
+@inline function _family_logd(d::Logistic, x)
+    z = (x - d.μ) / d.θ
+    -z - 2 * log1pexp(-z) - log(d.θ)
+end
+@inline _family_insupport(::Logistic, x) = true
+
+@inline function _family_logd(d::Cauchy, x)
+    z = (x - d.μ) / d.σ
+    -log1p(z * z) - log(π * d.σ)
+end
+@inline _family_insupport(::Cauchy, x) = true
+
+@inline _family_logd(d::Laplace, x) = -abs((x - d.μ) / d.θ) - log(2 * d.θ)
+@inline _family_insupport(::Laplace, x) = true
+
+@inline function _family_logd(d::LogNormal, x)
+    lx = log(abs(x))
+    z = (lx - d.μ) / d.σ
+    ℓ = -z * z / 2 - log(d.σ) - log2π / 2 - lx
+    ifelse(x > 0, ℓ, oftype(ℓ, -Inf))
+end
+@inline _family_insupport(::LogNormal, x) = x >= 0
+
+@inline function _family_logd(d::Weibull, x)
+    xθ = abs(x / d.θ)
+    ℓ = log(d.α / d.θ) + _clog(d.α - 1, xθ) - xθ^d.α
+    ifelse(isinf(xθ), oftype(ℓ, -Inf), ℓ)
+end
+@inline _family_insupport(::Weibull, x) = x >= 0
+
+@inline function _family_logd(d::Gamma, x)
+    ℓ = _clog(d.α - 1, abs(x)) - x / d.θ - loggamma(d.α) - d.α * log(d.θ)
+    ifelse(isinf(x), oftype(ℓ, -Inf), ℓ)
+end
+@inline _family_insupport(::Gamma, x) = x >= 0
+
+@inline function _family_logd(d::Beta, x)
+    _clog(d.α - 1, abs(x)) + _clog(d.β - 1, abs(1 - x)) - logbeta(d.α, d.β)
+end
+@inline _family_insupport(::Beta, x) = (0 <= x) & (x <= 1)
+
+@inline _family_logd(d::Poisson, x) = _clog(x, d.λ) - d.λ - loggamma(abs(x) + 1)
+@inline _family_insupport(::Poisson, x) = (x >= 0) & (x == floor(x))
+
+@inline _family_logd(d::Bernoulli, x) = ifelse(x == 1, log(d.p), log1p(-d.p))
+@inline _family_insupport(::Bernoulli, x) = (x == 0) | (x == 1)
+
+
+# Standard transports of the non-affine families: Cauchy, Laplace, Gamma
+# and Beta pivot on the uniform measure, the log-normal and Weibull
+# families on the normal resp. exponential measure.
+
+# The regularized incomplete gamma and beta functions and their inverses
+# (from SpecialFunctions), with derivatives with respect to the variate
+# resp. probability argument provided by the autodiff extensions:
+@inline MeasureBase._gamma_cdf(α, x) = MeasureBase._gamma_cdf_impl(MeasureBase._dualtag(α, x), α, x)
+@inline MeasureBase._gamma_quantile(α, p) = MeasureBase._gamma_quantile_impl(MeasureBase._dualtag(α, p), α, p)
+@inline MeasureBase._beta_cdf(α, β, x) = MeasureBase._beta_cdf_impl(MeasureBase._dualtag(α, β, x), α, β, x)
+@inline MeasureBase._beta_quantile(α, β, p) = MeasureBase._beta_quantile_impl(MeasureBase._dualtag(α, β, p), α, β, p)
+@inline MeasureBase._gamma_cdf_impl(::Type{Nothing}, α, x) = first(gamma_inc(α, x))
+# The complementary probability is formed in the common float type, as
+# `gamma_inc_inv` requires `p + q == 1` exactly:
+@inline function MeasureBase._gamma_quantile_impl(::Type{Nothing}, α, p)
+    T = float(promote_type(typeof(α), typeof(p)))
+    pp = convert(T, p)
+    gamma_inc_inv(convert(T, α), pp, one(T) - pp)
+end
+@inline MeasureBase._beta_cdf_impl(::Type{Nothing}, α, β, x) = first(beta_inc(α, β, x))
+@inline MeasureBase._beta_quantile_impl(::Type{Nothing}, α, β, p) = first(beta_inc_inv(α, β, p))
+@inline MeasureBase._gamma_logpdf(α, x) = _clog(α - 1, x) - x - loggamma(α)
+@inline MeasureBase._beta_logpdf(α, β, x) = _clog(α - 1, x) + _clog(β - 1, 1 - x) - logbeta(α, β)
+
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Cauchy}) = StdUniform
+@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Cauchy, x) = 1 // 2 + atan((x - d.μ) / d.σ) / π
+@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Cauchy, p) = muladd(d.σ, tan(π * (p - 1 // 2)), d.μ)
+
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Laplace}) = StdUniform
+@inline function MeasureBase.transport_to_std(::Type{StdUniform}, d::Laplace, x)
+    z = (x - d.μ) / d.θ
+    ifelse(z < 0, exp(z) / 2, 1 - exp(-z) / 2)
+end
+@inline function MeasureBase.transport_from_std(::Type{StdUniform}, d::Laplace, p)
+    u = p - 1 // 2
+    muladd(-d.θ * sign(u), log1p(-2 * abs(u)), d.μ)
+end
+
+@inline MeasureBase.preferred_stdmeasure(::Type{<:LogNormal}) = StdNormal
+@inline MeasureBase.transport_to_std(::Type{StdNormal}, d::LogNormal, x) = (log(x) - d.μ) / d.σ
+@inline MeasureBase.transport_from_std(::Type{StdNormal}, d::LogNormal, z) = exp(muladd(d.σ, z, d.μ))
+
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Weibull}) = StdExponential
+@inline MeasureBase.transport_to_std(::Type{StdExponential}, d::Weibull, x) = (x / d.θ)^d.α
+@inline MeasureBase.transport_from_std(::Type{StdExponential}, d::Weibull, z) = d.θ * z^(1 / d.α)
+
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Gamma}) = StdUniform
+@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Gamma, x) = _gamma_cdf(d.α, x / d.θ)
+@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Gamma, p) = d.θ * _gamma_quantile(d.α, p)
+
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Beta}) = StdUniform
+@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Beta, x) = _beta_cdf(d.α, d.β, x)
+@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Beta, p) = _beta_quantile(d.α, d.β, p)
diff --git a/ext/MeasureBaseDistributionsExt/multivariate.jl b/ext/MeasureBaseDistributionsExt/multivariate.jl
new file mode 100644
index 00000000..6ca640f9
--- /dev/null
+++ b/ext/MeasureBaseDistributionsExt/multivariate.jl
@@ -0,0 +1,112 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+# Multivariate normal and Dirichlet measures with density kernels and
+# transports over flat batches `(n, batch dims...)`, in terms of array
+# operations on the parameters, so that they run on devices and in traced
+# code. Single variates are vectors.
+
+# Column batches: flat batches as `(n, :)` matrices, single variates stay
+# vectors. Reductions over the columns give arrays over the batch
+# dimensions, numbers for single variates.
+@inline _as_columns(x::AbstractVector) = x
+@inline _as_columns(X::AbstractArray) = reshape(X, (size(X, 1), :))
+@inline _from_columns(y::AbstractVector, ::AbstractVector) = y
+@inline _from_columns(Y::AbstractMatrix, X::AbstractArray) = reshape(Y, (size(Y, 1), Base.tail(size(X))...))
+@inline _column_sums(f, z::AbstractVector) = sum(f, z)
+@inline _column_sums(f, Z::AbstractMatrix) = vec(sum(f, Z; dims = 1))
+@inline _column_all(z::AbstractVector) = all(z)
+@inline _column_all(Z::AbstractMatrix) = vec(all(Z; dims = 1))
+@inline _batch_results(r::Number, ::AbstractVector) = r
+@inline _batch_results(r::AbstractVector, X::AbstractArray) = reshape(r, Base.tail(size(X)))
+@inline _rows(z::AbstractVector, r) = view(z, r)
+@inline _rows(Z::AbstractMatrix, r) = view(Z, r, :)
+@inline _masked(ℓ::Number, ins) = ifelse(ins, ℓ, oftype(ℓ, -Inf))
+@inline _masked(ℓ::AbstractArray, ins) = ifelse.(ins, ℓ, eltype(ℓ)(-Inf))
+
+
+# Multivariate normal: densities via the Cholesky factor of the covariance.
+
+const MvNormalMeasure = AsMeasure{<:MvNormal}
+
+_logdet_cov(Σ::PDMats.PDMat) = 2 * sum(log, diag(Σ.chol.factors))
+_logdet_cov(Σ::PDMats.PDiagMat) = sum(log, Σ.diag)
+_logdet_cov(Σ::PDMats.ScalMat) = Σ.dim * log(Σ.value)
+
+for bhead in (:batched_logdensityof_impl, :batched_logdensity_def)
+    @eval function MeasureBase.$bhead(m::MvNormalMeasure, X::AbstractArray)
+        d = m.obj
+        Z = _cholesky_L(d.Σ) \ (_as_columns(X) .- d.μ)
+        sq = _column_sums(abs2, Z)
+        _batch_results(-sq ./ 2 .- (_logdet_cov(d.Σ) + length(d) * log2π) / 2, X)
+    end
+end
+MeasureBase.logdensity_def(m::MvNormalMeasure, x::AbstractVector) = MeasureBase.batched_logdensity_def(m, x)
+MeasureBase.unsafe_logdensityof(m::MvNormalMeasure, x::AbstractVector) = MeasureBase.batched_logdensityof_impl(m, x)
+
+function MeasureBase.batched_transport_to_std(::Type{StdNormal}, d::MvNormal, X::AbstractArray)
+    _from_columns(_cholesky_L(d.Σ) \ (_as_columns(X) .- d.μ), X)
+end
+function MeasureBase.batched_transport_from_std(::Type{StdNormal}, d::MvNormal, Z::AbstractArray)
+    _from_columns(_cholesky_L(d.Σ) * _as_columns(Z) .+ d.μ, Z)
+end
+MeasureBase.transport_to_std(::Type{StdNormal}, d::MvNormal, x) = MeasureBase.batched_transport_to_std(StdNormal, d, x)
+MeasureBase.transport_from_std(::Type{StdNormal}, d::MvNormal, z) = MeasureBase.batched_transport_from_std(StdNormal, d, z)
+
+# Parameters follow the batches to the device:
+function Adapt.adapt_structure(to, m::MvNormalMeasure)
+    d = m.obj
+    asmeasure(MvNormal(Adapt.adapt(to, d.μ), _adapt_cov(to, d.Σ)))
+end
+function _adapt_cov(to, Σ::PDMats.PDMat)
+    chol = Σ.chol
+    PDMats.PDMat(Adapt.adapt(to, Σ.mat), Cholesky(Adapt.adapt(to, chol.factors), chol.uplo, chol.info))
+end
+_adapt_cov(to, Σ::PDMats.PDiagMat) = PDMats.PDiagMat(Adapt.adapt(to, Σ.diag))
+_adapt_cov(to, Σ::PDMats.ScalMat) = Σ
+
+
+# Dirichlet: densities over the simplex, transports via the stick-breaking
+# Beta transports (M. J. Betancourt, "Cruising The Simplex: Hamiltonian
+# Monte Carlo and the Dirichlet Distribution", arXiv:1010.3436), with the
+# cumulative sums and products running along the variate dimension.
+
+for bhead in (:batched_logdensityof_impl, :batched_logdensity_def)
+    @eval function MeasureBase.$bhead(m::DirichletMeasure, X::AbstractArray)
+        d = m.obj
+        Xc = _as_columns(X)
+        ℓ = _column_sums(identity, _clog.(d.alpha .- 1, abs.(Xc))) .- d.lmnB
+        tol = sqrt(eps(float(eltype(X))))
+        ins = _column_all(Xc .>= 0) .& (abs.(_column_sums(identity, Xc) .- 1) .<= tol)
+        _batch_results(_masked(ℓ, ins), X)
+    end
+end
+MeasureBase.logdensity_def(m::DirichletMeasure, x::AbstractVector) = MeasureBase.batched_logdensity_def(m, x)
+MeasureBase.unsafe_logdensityof(m::DirichletMeasure, x::AbstractVector) = MeasureBase.batched_logdensityof_impl(m, x)
+
+# The stick-breaking Beta parameters, for the first `K - 1` components:
+@inline _stick_breaking_params(d::Dirichlet) = (_dropfront(_rev_cumsum(d.alpha)), _dropback(d.alpha))
+
+function MeasureBase.batched_transport_to_std(::Type{StdUniform}, d::Dirichlet, X::AbstractArray)
+    K = length(d)
+    αs, βs = _stick_breaking_params(d)
+    Xc = _as_columns(X)
+    rem = 1 .- cumsum(Xc; dims = 1)
+    # The remaining mass before each component is the mass after it plus
+    # the component itself:
+    beta_v = _rows(rem, 1:(K - 1)) ./ (_rows(rem, 1:(K - 1)) .+ _rows(Xc, 1:(K - 1)))
+    _from_columns(_beta_cdf.(αs, βs, beta_v), X)
+end
+
+function MeasureBase.batched_transport_from_std(::Type{StdUniform}, d::Dirichlet, Z::AbstractArray)
+    K = length(d)
+    αs, βs = _stick_breaking_params(d)
+    beta_v = _beta_quantile.(αs, βs, _as_columns(Z))
+    cp = cumprod(beta_v; dims = 1)
+    # Each component takes what its Beta variate leaves of the remaining mass:
+    X = vcat(1 .- _rows(cp, 1:1), _rows(cp, 1:(K - 2)) .- _rows(cp, 2:(K - 1)), _rows(cp, (K - 1):(K - 1)))
+    _from_columns(X, Z)
+end
+MeasureBase.transport_to_std(::Type{StdUniform}, d::Dirichlet, x) = MeasureBase.batched_transport_to_std(StdUniform, d, x)
+MeasureBase.transport_from_std(::Type{StdUniform}, d::Dirichlet, z) = MeasureBase.batched_transport_from_std(StdUniform, d, z)
+
+Adapt.adapt_structure(to, m::DirichletMeasure) = asmeasure(Dirichlet(Adapt.adapt(to, m.obj.alpha)))
diff --git a/ext/MeasureBaseDistributionsExt/standardmv.jl b/ext/MeasureBaseDistributionsExt/standardmv.jl
index 1163c39b..1c9293aa 100644
--- a/ext/MeasureBaseDistributionsExt/standardmv.jl
+++ b/ext/MeasureBaseDistributionsExt/standardmv.jl
@@ -6,32 +6,13 @@ MeasureBase.getdof(m::AsMeasure{<:AbstractMvNormal}) = getdof(m.obj)
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:AbstractMvNormal}) = StdNormal
 
-_cholesky_L(A) = cholesky(A).L
+_cholesky_L(A) = _lower_factor(cholesky(A))
+# `Cholesky.L` copies the transposed factor and scalar-indexes on GPUs.
+_lower_factor(C::Cholesky) = C.uplo === 'L' ? LowerTriangular(C.factors) : UpperTriangular(C.factors)'
 _cholesky_L(A::Diagonal{<:Real}) = Diagonal(sqrt.(diag(A)))
 _cholesky_L(A::PDMats.PDiagMat{<:Real}) = Diagonal(sqrt.(A.diag))
 _cholesky_L(A::PDMats.ScalMat{<:Real}) = Diagonal(Fill(sqrt(A.value), A.dim))
 
-function MeasureBase.transport_to_std(::Type{StdNormal}, d::MvNormal, x)
-    _cholesky_L(d.Σ) \ (x - d.μ)
-end
-
-function MeasureBase.transport_from_std(::Type{StdNormal}, d::MvNormal, z)
-    muladd(_cholesky_L(d.Σ), z, d.μ)
-end
-
-function MeasureBase.batched_transport_to_std(::Type{StdNormal}, d::MvNormal, X::AbstractArray)
-    X_mat = reshape(X, (length(d), :))
-    Z_mat = _cholesky_L(d.Σ) \ (X_mat .- d.μ)
-    return reshape(Z_mat, size(X))
-end
-
-function MeasureBase.batched_transport_from_std(::Type{StdNormal}, d::MvNormal, Z::AbstractArray)
-    Z_mat = reshape(Z, (length(d), :))
-    X_mat = muladd(_cholesky_L(d.Σ), Z_mat, d.μ)
-    return reshape(X_mat, size(Z))
-end
-
-
 #DirichletMultinomial
 #Distributions.AbstractMvLogNormal
 #Distributions.AbstractMvTDist
diff --git a/ext/MeasureBaseDistributionsForwardDiffExt.jl b/ext/MeasureBaseDistributionsForwardDiffExt.jl
index e929066f..181fc410 100644
--- a/ext/MeasureBaseDistributionsForwardDiffExt.jl
+++ b/ext/MeasureBaseDistributionsForwardDiffExt.jl
@@ -57,6 +57,50 @@ end
     ForwardDiff.Dual{TAG}(x, dx_dp * ForwardDiff.partials(p))
 end
 
+# Dual numbers through the regularized incomplete gamma and beta functions
+# and their inverses: the derivative with respect to the variate resp.
+# probability argument follows from the density, derivatives with respect
+# to the parameters are not available and yield NaN partials.
+using MeasureBase: _gamma_cdf, _gamma_quantile, _beta_cdf, _beta_quantile, _gamma_logpdf, _beta_logpdf
+
+const _Dual = ForwardDiff.Dual
+
+@inline MeasureBase._dualtag(::_Dual{TAG}, ::Number...) where {TAG} = _Dual{TAG}
+
+# The derivative through the last argument; the partials of parameters
+# are marked NaN where they are nonzero (their derivatives are unknown):
+@inline _arg_partials(::Type{_Dual{TAG}}, x::_Dual{TAG}, params...) where {TAG} = ForwardDiff.partials(x)
+@inline _arg_partials(::Type{_Dual{TAG}}, ::Real, params...) where {TAG} = zero(ForwardDiff.partials(_first_dual(params...)))
+@inline _first_dual(x::_Dual, rest...) = x
+@inline _first_dual(::Real, rest...) = _first_dual(rest...)
+@inline _nan_partials(∂, ::Real) = ∂
+@inline function _nan_partials(∂, x::_Dual)
+    ∂ + ForwardDiff.Partials(map(v -> ifelse(iszero(v), zero(v), oftype(v, NaN)), ForwardDiff.partials(x).values))
+end
+@inline function _through_last(::Type{_Dual{TAG}}, value, dvalue, last, params...) where {TAG}
+    ∂ = dvalue * _arg_partials(_Dual{TAG}, last, params...)
+    ForwardDiff.Dual{TAG}(value, foldl(_nan_partials, params; init = ∂))
+end
+
+@inline function MeasureBase._gamma_cdf_impl(::Type{_Dual{TAG}}, α, x) where {TAG}
+    αv, xv = ForwardDiff.value(α), ForwardDiff.value(x)
+    _through_last(_Dual{TAG}, _gamma_cdf(αv, xv), exp(_gamma_logpdf(αv, xv)), x, α)
+end
+@inline function MeasureBase._gamma_quantile_impl(::Type{_Dual{TAG}}, α, p) where {TAG}
+    αv, pv = ForwardDiff.value(α), ForwardDiff.value(p)
+    xv = _gamma_quantile(αv, pv)
+    _through_last(_Dual{TAG}, xv, exp(-_gamma_logpdf(αv, xv)), p, α)
+end
+@inline function MeasureBase._beta_cdf_impl(::Type{_Dual{TAG}}, α, β, x) where {TAG}
+    αv, βv, xv = ForwardDiff.value(α), ForwardDiff.value(β), ForwardDiff.value(x)
+    _through_last(_Dual{TAG}, _beta_cdf(αv, βv, xv), exp(_beta_logpdf(αv, βv, xv)), x, α, β)
+end
+@inline function MeasureBase._beta_quantile_impl(::Type{_Dual{TAG}}, α, β, p) where {TAG}
+    αv, βv, pv = ForwardDiff.value(α), ForwardDiff.value(β), ForwardDiff.value(p)
+    xv = _beta_quantile(αv, βv, pv)
+    _through_last(_Dual{TAG}, xv, exp(-_beta_logpdf(αv, βv, xv)), p, α, β)
+end
+
 # The quantile of Beta doesn't support dual parameters:
 @inline MeasureBase._dist_quantile(d::Beta{<:ForwardDiff.Dual}, p::Real) = convert(float(typeof(p)), NaN)
 @inline MeasureBase._dist_cquantile(d::Beta{<:ForwardDiff.Dual}, p::Real) = convert(float(typeof(p)), NaN)
diff --git a/redesign.md b/redesign.md
index 333260be..518c32e0 100644
--- a/redesign.md
+++ b/redesign.md
@@ -122,19 +122,29 @@ dependency, `MeasureBaseAffineMapsExt`), a host loop otherwise.
 - Weighted, restricted, density measures, Half: forward plus weights or
   masks. Superpositions and spike mixtures: one batch per component,
   masks aligned with the variate dimensions. Dirac: constant batches.
-- Distributions extension: univariate via `StdLogistic` (log-cdf and
-  quantile), location-scale families via their affine map, `MvNormal` via
-  Cholesky factors, array-variate batches via `logpdf(d, X)`.
+- Distributions extension: the main univariate families (Normal,
+  Uniform, Exponential, Logistic, Cauchy, Laplace, LogNormal, Weibull,
+  Gamma, Beta, Poisson, Bernoulli) get branch-free density formulas and
+  transports to the standard measure matching their shape (`families.jl`),
+  other univariate distributions go via `StdLogistic` (log-cdf and
+  quantile). `MvNormal` works on column batches through its Cholesky
+  factor, `Dirichlet` through stick-breaking Beta transports with `cumsum`
+  and `cumprod` along the variate dimension (`multivariate.jl`). Incomplete
+  gamma and beta functions go through the `_gamma_cdf`/`_beta_cdf` hooks
+  (and quantiles), whose ForwardDiff and ChainRules derivatives come from
+  the densities. Draws use the Distributions samplers on the CPU and the
+  standard transports on other compute units. Remaining array-variate
+  distributions batch via `logpdf(d, X)` on the host.
 
 ## Changes relative to `master`
 
-Twenty commits since `c773afe`: variate size contract and
+Commits since `c773afe`: variate size contract and
 `preferred_stdmeasure`, densities over flat storage, branch-free
 evaluation, Reactant smoke tests, structural batched kernels, transport
 rebuilt on standard measures, batched transport and the broadcast hook,
 rand via `GenContext`, then the batched-first redesign (density core,
-struct array products, transport, random variates, structured batches)
-and the review fixes.
+struct array products, transport, random variates, structured batches),
+the review fixes, and device-friendly kernels for wrapped Distributions.
 
 Removed: `transport_origin`/`to_origin`/`from_origin` and the origin
 machinery, `NoTransport`, `transport_to_mvstd`, per-measure
@@ -152,18 +162,25 @@ products gives plain arrays); vcat-combined and bind variates are flat;
 
 ## Verification
 
-Full suite (Aqua, extensions, doctests) on CPU with JLArrays cases;
-`test/reactant` (opt-in, backend via `MEASUREBASE_REACTANT_BACKEND`) and
-`test/cuda` (opt-in) run locally on the GB10, both green at HEAD except
-one expected-broken CUDA case (AffineMaps Jacobian on device).
+Full suite (Aqua, extensions, doctests) on CPU with JLArrays cases.
+`test/test_reactant.jl` runs as part of the suite on 64-bit Linux and
+macOS with stable Julia, adding Reactant on demand as MGVI does (backend
+via `MEASUREBASE_REACTANT_BACKEND`). `test/cuda` is opt-in. Both run
+locally on the GB10, green at HEAD except one expected-broken CUDA case
+(AffineMaps Jacobian on device).
 
 ## Known gaps and open decisions
 
 - Upstream: AffineMaps lacks `Adapt` rules and device/traced Jacobians;
-  Distributions isn't device-aware; ChangesOfVariables has no rules for
-  `Base.Fix1`/`Fix2` arithmetic; HeterogeneousComputing has no Reactant
-  compute unit; JLArrays has no RNG; Reactant rejects traced
-  `VectorOfArrays` and empty batches.
+  Distributions isn't device-aware and its parameter structs (PDMats,
+  `Dirichlet`) can't hold traced arrays, so under Reactant distribution
+  parameters stay constants; SpecialFunctions' incomplete beta and gamma
+  functions don't compile for GPUs or Reactant, so Beta and Dirichlet
+  transports and draws (Gamma under Reactant too) stay on the CPU;
+  ChangesOfVariables has no rules for `Base.Fix1`/`Fix2` arithmetic;
+  HeterogeneousComputing has no Reactant compute unit; JLArrays has no
+  RNG and no triangular solves; Reactant rejects traced `VectorOfArrays`
+  and empty batches.
 - Decisions pending: one convention for out-of-support inputs (NaN mask
   vs. DomainError vs. AssertionError), `Half` tails via log-ccdf, device
   random variate infrastructure and `rand!`, Tier-1 static variates,
diff --git a/src/utils.jl b/src/utils.jl
index cd215682..397c379a 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -183,6 +183,27 @@ convert_realtype(::Type{T}, x::Union{Tuple,NamedTuple}) where {T<:Real} =
 convert_realtype(::Type{T}, x::AbstractArray) where {T<:Real} =
     map(Base.Fix1(convert_realtype, T), x)
 
+# Regularized incomplete gamma and beta functions and their inverses, with
+# the log-densities of the standard gamma and beta distributions for their
+# derivatives. Implemented in the Distributions extension, differentiated
+# with respect to the variate resp. probability argument in the autodiff
+# extensions:
+function _gamma_cdf end
+function _gamma_quantile end
+function _beta_cdf end
+function _beta_quantile end
+function _gamma_logpdf end
+function _beta_logpdf end
+function _gamma_cdf_impl end
+function _gamma_quantile_impl end
+function _beta_cdf_impl end
+function _beta_quantile_impl end
+
+# The dual number type among the arguments of such a function, `Nothing`
+# for plain numbers (the ForwardDiff extension adds dual numbers):
+@inline _dualtag() = Nothing
+@inline _dualtag(::Number, rest::Number...) = _dualtag(rest...)
+
 # Distributions implementation hooks, specialized for dual numbers in the
 # ForwardDiff extension:
 function _trafo_logcdf_impl end
diff --git a/test/cuda/runtests.jl b/test/cuda/runtests.jl
index dd34cd79..aaaba760 100644
--- a/test/cuda/runtests.jl
+++ b/test/cuda/runtests.jl
@@ -9,13 +9,14 @@ using CUDA
 using Adapt: adapt
 using HeterogeneousComputing: GenContext, AbstractComputeUnit
 using MeasureBase
-using MeasureBase: StdNormal, StdUniform, StdExponential, Dirac
+using MeasureBase: StdNormal, StdUniform, StdExponential, Dirac, asmeasure
 using MeasureBase: productmeasure, pushfwd, mcombine, weightedmeasure, superpose, SpikeMixture
 using MeasureBase: batched_rand_impl
 using MeasureBase.InverseFunctions: inverse
 using ArraysOfArrays: sliced, flatview
 using AffineMaps: Mul, MulAdd
-using Distributions: Normal
+using Distributions: Normal, Uniform, Exponential, Logistic, Cauchy, Laplace, LogNormal, Weibull, Gamma, Beta
+using Distributions: Poisson, Bernoulli, MvNormal, Dirichlet
 
 CUDA.allowscalar(false)
 
@@ -100,9 +101,39 @@ _plain(x::Tuple) = map(_plain, x)
         @test_broken Array(logdensities(νac, CuArray(Ya))) ≈ logdensities(νa, Ya)
     end
 
+    @testset "wrapped distributions" begin
+        for d in (Normal(0.3, 1.7), Uniform(-1.0, 2.5), Exponential(0.7), Logistic(0.2, 1.3), Cauchy(0.1, 0.8), Laplace(-0.4, 1.1), LogNormal(0.2, 0.6), Weibull(1.4, 0.9), Gamma(2.3, 1.2), Beta(2.5, 3.5))
+            m = asmeasure(d)
+            xd = rand(d, 20)
+            test_cuda(x -> logdensities(m, x), xd)
+            f = transport_to(StdNormal(), m)
+            if d isa Beta
+                # SpecialFunctions' incomplete beta function doesn't compile for GPUs:
+                @test_broken Array(f.(CuArray(xd))) ≈ f.(xd)
+            else
+                test_cuda(x -> f.(x), xd)
+                test_cuda(z -> inverse(f).(z), randn(20))
+            end
+        end
+        for d in (Poisson(2.7), Bernoulli(0.3))
+            test_cuda(x -> logdensities(asmeasure(d), x), Float64.(rand(d, 20)))
+        end
+        mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3])
+        mm = asmeasure(mvn)
+        Xm = rand(mvn, 20)
+        test_cuda((m, X) -> logdensities(m, X), mm, Xm)
+        test_cuda((m, X) -> flatview(transport_to(StdNormal()^2, m).(sliced(X, Val(1)))), mm, Xm)
+        test_cuda((m, Z) -> flatview(transport_to(m, StdNormal()^2).(sliced(Z, Val(1)))), mm, randn(2, 20))
+        dir = Dirichlet([2.0, 3.0, 4.0, 1.5])
+        md = asmeasure(dir)
+        Xd = rand(dir, 20)
+        test_cuda((m, X) -> logdensities(m, X), md, Xd)
+        @test_broken flatview(transport_to(StdUniform()^3, cu_copy(md)).(sliced(CuArray(Xd), Val(1)))) isa CuArray
+    end
+
     @testset "random variates" begin
         ctx = GenContext{Float32}(AbstractComputeUnit(CUDA.device()), CUDA.default_rng())
-        for μ in (StdNormal(), StdUniform(), StdExponential(), StdNormal()^3, (StdNormal()^2)^3, weightedmeasure(0.3, StdNormal()^2), mcombine(vcat, StdNormal()^2, StdUniform()^1), pushfwd(Base.BroadcastFunction(exp), StdNormal()^2), superpose(StdNormal(), StdUniform()), SpikeMixture(StdNormal(), 0.5))
+        for μ in (StdNormal(), StdUniform(), StdExponential(), StdNormal()^3, (StdNormal()^2)^3, weightedmeasure(0.3, StdNormal()^2), mcombine(vcat, StdNormal()^2, StdUniform()^1), pushfwd(Base.BroadcastFunction(exp), StdNormal()^2), superpose(StdNormal(), StdUniform()), SpikeMixture(StdNormal(), 0.5), asmeasure(Normal(0.3, 1.7)), asmeasure(Weibull(1.4, 0.9)), cu_copy(asmeasure(MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3]))))
             X = batched_rand_impl(ctx, μ, (100,))
             @test X isa CuArray{Float32}
             ℓ = logdensities(μ, X)
diff --git a/test/distributions/test_device_kernels.jl b/test/distributions/test_device_kernels.jl
new file mode 100644
index 00000000..c4b53932
--- /dev/null
+++ b/test/distributions/test_device_kernels.jl
@@ -0,0 +1,114 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+# The device-friendly kernels of the wrapped distribution families: their
+# densities and transports agree with Distributions, and they run on flat
+# batches, also of device arrays.
+
+using Test
+using Distributions, LinearAlgebra, StableRNGs, Statistics
+using MeasureBase
+using MeasureBase: asmeasure, GenContext, StdUniform, StdNormal, StdExponential, StdLogistic
+using MeasureBase: batched_rand_impl, logdensities, insupport
+using MeasureBase.InverseFunctions: inverse
+using ArraysOfArrays: sliced, flatview
+import Adapt
+using JLArrays
+
+@testset "device kernels of distribution families" begin
+    JLArrays.allowscalar(false)
+    stblrng() = StableRNG(28734)
+
+    families = [
+        Normal(0.3, 1.7), Uniform(-1.0, 2.5), Exponential(0.7), Logistic(0.2, 1.3), Cauchy(0.1, 0.8),
+        Laplace(-0.4, 1.1), LogNormal(0.2, 0.6), Weibull(1.4, 0.9), Gamma(2.3, 1.2), Beta(2.5, 3.5),
+        Normal(0.3f0, 1.7f0), Gamma(0.7, 2.0), Beta(0.6, 0.8),
+    ]
+
+    @testset "$(nameof(typeof(d)))" for d in families
+        m = asmeasure(d)
+        xs = vcat(rand(stblrng(), d, 20), [-1.0, 0.0, 1.0, 5.0, Inf])
+        xs = eltype(d) == Float32 ? Float32.(xs) : xs
+        ℓ_ref = logpdf.(d, xs)
+        @test all(map((a, b) -> a == b || a ≈ b || (isnan(a) && isnan(b)), logdensityof.(Ref(m), xs), ℓ_ref))
+        @test logdensities(m, xs) ≈ ℓ_ref nans = true
+        @test Array(logdensities(m, JLArray(xs))) ≈ ℓ_ref nans = true
+        @test insupport.(Ref(m), xs) == Distributions.insupport.(d, xs)
+
+        if d isa ContinuousUnivariateDistribution
+            x = rand(stblrng(), d, 12)
+            f = transport_to(StdUniform(), m)
+            p = f.(x)
+            @test p ≈ cdf.(d, x)
+            @test inverse(f).(p) ≈ x
+            @test Array(inverse(f).(f.(JLArray(x)))) ≈ x
+            g = transport_to(StdNormal(), m)
+            @test inverse(g).(g.(x)) ≈ x
+            X = batched_rand_impl(GenContext{Float64}(stblrng()), m, (5,))
+            @test X isa Vector{Float64} && length(X) == 5
+            # The transport-based generation used on devices, in single precision:
+            Xs = MeasureBase._rand_default(GenContext{Float32}(stblrng()), m, (2000,), MeasureBase._NoRandImpl())
+            @test Xs isa Vector{Float32} && all(insupport.(Ref(m), Xs))
+            if isfinite(mean(d)) && isfinite(var(d))
+                @test isapprox(mean(Xs), mean(d), atol = 5 * sqrt(var(d) / 2000) + 1e-3)
+            end
+        end
+    end
+
+    @testset "discrete families" begin
+        for d in (Poisson(2.7), Bernoulli(0.3))
+            m = asmeasure(d)
+            xs = [0, 1, 2, 3, 7]
+            @test logdensityof.(Ref(m), xs) ≈ logpdf.(d, xs)
+            @test logdensities(m, xs) ≈ logpdf.(d, xs)
+            @test Array(logdensities(m, JLArray(xs))) ≈ logpdf.(d, xs)
+            @test logdensityof(m, -1) == -Inf && logdensityof(m, 1.5) == -Inf
+        end
+    end
+
+    @testset "MvNormal" begin
+        for Σ in [[1.7 0.5; 0.5 2.3], PDMats.PDiagMat([0.5, 2.0]), PDMats.ScalMat(2, 1.5)]
+            d = MvNormal([0.3, -2.9], Σ)
+            m = asmeasure(d)
+            X = rand(stblrng(), d, 6)
+            ℓ_ref = logpdf(d, X)
+            @test logdensityof(m, X[:, 1]) ≈ ℓ_ref[1]
+            @test logdensities(m, X) ≈ ℓ_ref
+            @test logdensities(m, sliced(X, Val(1))) ≈ ℓ_ref
+            @test logdensityof(m^6, X) ≈ sum(ℓ_ref)
+            f = transport_to(StdNormal()^2, m)
+            Y = f.(sliced(X, Val(1)))
+            @test flatview(Y) ≈ stack(map(f, eachcol(X)))
+            @test flatview(inverse(f).(Y)) ≈ X
+            # JLArrays have no triangular solves, so only diagonal
+            # covariances run on them (CUDA covers the general case):
+            if !(Σ isa AbstractMatrix)
+                mj = Adapt.adapt(JLArray, m)
+                @test Array(logdensities(mj, JLArray(X))) ≈ ℓ_ref
+                fj = transport_to(StdNormal()^2, mj)
+                @test Array(flatview(fj.(sliced(JLArray(X), Val(1))))) ≈ flatview(Y)
+                @test Array(flatview(inverse(fj).(sliced(JLArray(flatview(Y)), Val(1))))) ≈ X
+            end
+            @test size(batched_rand_impl(GenContext{Float64}(stblrng()), m, (7,))) == (2, 7)
+        end
+    end
+
+    @testset "Dirichlet" begin
+        d = Dirichlet([2.0, 3.0, 4.0, 1.5])
+        m = asmeasure(d)
+        X = rand(stblrng(), d, 6)
+        ℓ_ref = logpdf(d, X)
+        @test logdensityof(m, X[:, 1]) ≈ ℓ_ref[1]
+        @test logdensities(m, X) ≈ ℓ_ref
+        @test logdensityof(m, [0.5, 0.5, 0.2, 0.1]) == -Inf
+        @test logdensityof(m, [0.5, 0.6, -0.1, 0.0]) == -Inf
+        mj = Adapt.adapt(JLArray, m)
+        @test Array(logdensities(mj, JLArray(X))) ≈ ℓ_ref
+        f = transport_to(StdUniform()^3, m)
+        Y = f.(sliced(X, Val(1)))
+        @test flatview(Y) ≈ stack(map(f, eachcol(X)))
+        @test flatview(inverse(f).(Y)) ≈ X
+        @test all(0 .<= flatview(Y) .<= 1)
+        Xr = batched_rand_impl(GenContext{Float64}(stblrng()), m, (200,))
+        @test size(Xr) == (4, 200) && all(sum(Xr; dims = 1) .≈ 1)
+    end
+end
diff --git a/test/distributions/test_distributions.jl b/test/distributions/test_distributions.jl
index c73d50a7..9e5777b0 100644
--- a/test/distributions/test_distributions.jl
+++ b/test/distributions/test_distributions.jl
@@ -21,5 +21,6 @@ using .MeasureBaseDistributionsExt:
     include("test_standard_normal.jl")
     include("test_conversions.jl")
     include("test_transport.jl")
+    include("test_device_kernels.jl")
     include("test_mooncake.jl")
 end
diff --git a/test/distributions/test_shape_contract.jl b/test/distributions/test_shape_contract.jl
index 52d5b608..d4865fc1 100644
--- a/test/distributions/test_shape_contract.jl
+++ b/test/distributions/test_shape_contract.jl
@@ -19,7 +19,7 @@ using LinearAlgebra: I
     @test @inferred(preferred_stdmeasure(Uniform(1, 2))) === StdUniform
     @test @inferred(preferred_stdmeasure(Exponential(2.0))) === StdExponential
     @test @inferred(preferred_stdmeasure(Logistic(1, 2))) === StdLogistic
-    @test @inferred(preferred_stdmeasure(Beta(2, 3))) === StdLogistic
+    @test @inferred(preferred_stdmeasure(Beta(2, 3))) === StdUniform
     @test @inferred(preferred_stdmeasure(truncated(Normal(), 0, 1))) === StdLogistic
     @test @inferred(preferred_stdmeasure(MvNormal(zeros(2), I(2)))) === StdNormal
     @test @inferred(preferred_stdmeasure(Dirichlet([1.0, 2.0]))) === StdUniform
@@ -28,10 +28,10 @@ using LinearAlgebra: I
     @test @inferred(preferred_stdmeasure(StandardDist{Uniform}())) === StdUniform
 
     @test @inferred(preferred_stdmeasure(productmeasure((asmeasure(Beta(2, 3)), asmeasure(Normal()))))) === StdNormal
-    @test @inferred(preferred_stdmeasure(productmeasure((a = Dirac(1.0), b = asmeasure(Beta(2, 3)))))) === StdLogistic
+    @test @inferred(preferred_stdmeasure(productmeasure((a = Dirac(1.0), b = asmeasure(Beta(2, 3)))))) === StdUniform
     @test @inferred(preferred_stdmeasure(productmeasure((a = asmeasure(Poisson(2)), b = asmeasure(Beta(2, 3)))))) <: NoStdTransport
     @test @inferred(preferred_stdmeasure(productmeasure([asmeasure(Normal(i, 1)) for i in 1:3]))) === StdNormal
-    @test @inferred(preferred_stdmeasure(product_distribution([Beta(2, 3), Beta(1, 1)]))) === StdLogistic
+    @test @inferred(preferred_stdmeasure(product_distribution([Beta(2, 3), Beta(1, 1)]))) === StdUniform
 
     lkj = asmeasure(LKJCholesky(3, 1.0))
     @test @inferred(mspace_elsize(lkj)) isa MeasureBase.NoMSpaceElementSize

From 4559e8280de45df3c8354f91701a4a026ce28f86 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 14:40:57 +0200
Subject: [PATCH 102/122] Fix test issues found by CI

Pkg is a test dependency, the on-demand Reactant install needs it. The
allocation tests measure inside a function, since @allocated at top
level reports a boxed result on Julia 1.10.

Created by generative AI.
---
 test/Project.toml    |  1 +
 test/logdensities.jl | 11 +++++++----
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/test/Project.toml b/test/Project.toml
index 7c717b0d..e4ed8cc0 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -22,6 +22,7 @@ LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 OneTwoMany = "762dc654-8631-413a-a342-372a7419ad9d"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PropertyFunctions = "09e99361-2bb8-48a2-a80f-de58f0739eb4"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
diff --git a/test/logdensities.jl b/test/logdensities.jl
index cfa32e36..95ee7081 100644
--- a/test/logdensities.jl
+++ b/test/logdensities.jl
@@ -24,6 +24,9 @@ MeasureBase.basemeasure(::VecTestMeasure) = LebesgueBase()^2
 MeasureBase.insupport(::VecTestMeasure, x) = true
 MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s)
 
+# `@allocated` at top level reports a boxed result on Julia 1.10:
+_allocated(f::F, args::Vararg{Any,N}) where {F,N} = @allocated f(args...)
+
 @testset "logdensities" begin
     @testset "scalar variates" begin
         X = randn(10)
@@ -125,10 +128,10 @@ MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s)
         m3 = StdNormal()^static(3)
         xs = @SVector randn(3)
         @test @inferred(logdensityof(m3, xs)) ≈ sum(stdnormal_ld, xs)
-        @test @allocated(logdensityof(m3, xs)) == 0
+        @test _allocated(logdensityof, m3, xs) == 0
         xd = randn(3)
         @test @inferred(logdensityof(StdNormal()^3, xd)) ≈ sum(stdnormal_ld, xd)
-        @test @allocated(logdensityof(StdNormal()^3, xd)) == 0
+        @test _allocated(logdensityof, StdNormal()^3, xd) == 0
         Xs = @SMatrix randn(3, 4)
         @test @inferred(logdensities(m3, Xs)) ≈ vec(sum(stdnormal_ld.(Xs), dims = 1))
         @test logdensities(m3, Xs) isa SVector{4}
@@ -148,7 +151,7 @@ MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s)
         @test @inferred(logdensities(w, X)) ≈ [logdensityof(w, x) for x in eachcol(X)]
         xw = randn(3)
         @test @inferred(logdensityof(w, xw)) ≈ log(0.3) + sum(stdnormal_ld, xw)
-        @test @allocated(logdensityof(w, xw)) == 0
+        @test _allocated(logdensityof, w, xw) == 0
 
         ms = [weightedmeasure(log(i), StdNormal()) for i in 1:4]
         prod4 = productmeasure(ms)
@@ -156,7 +159,7 @@ MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s)
         @test @inferred(MeasureBase.mspace_elsize(prod4)) == (4,)
         xp = randn(4)
         @test @inferred(logdensityof(prod4, xp)) ≈ sum(log(i) + stdnormal_ld(xp[i]) for i in 1:4)
-        @test @allocated(logdensityof(prod4, xp)) == 0
+        @test _allocated(logdensityof, prod4, xp) == 0
         Xp = randn(4, 7)
         @test @inferred(logdensities(prod4, Xp)) ≈ [logdensityof(prod4, x) for x in eachcol(Xp)]
         @test logdensities(prod4, sliced(Xp, 1)) ≈ logdensities(prod4, Xp)

From 332970cd94bb011ab34eb712aac387a6c05df84d Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 14:53:59 +0200
Subject: [PATCH 103/122] Measure test allocations through a shared helper

All allocation checks go through allocations_of in test/testutils.jl,
which warms up and measures inside a function, since @allocated at top
level reports a boxed result on Julia 1.10.

Created by generative AI.
---
 test/batched_regressions.jl |  6 ++++--
 test/logdensities.jl        | 11 +++++------
 test/testutils.jl           |  9 +++++++++
 test/transport.jl           |  6 ++++--
 4 files changed, 22 insertions(+), 10 deletions(-)
 create mode 100644 test/testutils.jl

diff --git a/test/batched_regressions.jl b/test/batched_regressions.jl
index 0456dc58..8471dfe3 100644
--- a/test/batched_regressions.jl
+++ b/test/batched_regressions.jl
@@ -29,6 +29,8 @@ end
 MeasureBase.InverseFunctions.inverse(f::Scale) = Scale(inv(f.s))
 MeasureBase.ChangesOfVariables.with_logabsdet_jacobian(f::Scale, x) = (f(x), log(abs(f.s)))
 
+include("testutils.jl")
+
 @testset "batched regressions" begin
     @testset "products of function-wrapper marginals" begin
         pm = productmeasure([pushfwd(Scale(s), StdExponential()) for s in 0.1:0.2:0.9])
@@ -111,10 +113,10 @@ MeasureBase.ChangesOfVariables.with_logabsdet_jacobian(f::Scale, x) = (f(x), log
         m1 = mcombine(vcat, StdNormal(), StdExponential()^static(2))
         x1 = SVector(0.1, 0.2, 0.3)
         @test logdensityof(m1, x1) ≈ logdensityof(StdNormal(), 0.1) + logdensityof(StdExponential()^2, [0.2, 0.3])
-        @test @allocated(logdensityof(m1, x1)) == 0
+        @test allocations_of(logdensityof, m1, x1) == 0
         m2 = mcombine(vcat, StdNormal()^static(2), StdExponential()^static(3))
         x2 = SVector(0.1, 0.2, 0.3, 0.4, 0.5)
-        @test @allocated(logdensityof(m2, x2)) == 0
+        @test allocations_of(logdensityof, m2, x2) == 0
         g = transport_to(StdUniform()^3, StdNormal()^3)
         v = randn(3)
         @test g.(SVector{3}(v))[] ≈ g(v)
diff --git a/test/logdensities.jl b/test/logdensities.jl
index 95ee7081..d1c8a01d 100644
--- a/test/logdensities.jl
+++ b/test/logdensities.jl
@@ -24,8 +24,7 @@ MeasureBase.basemeasure(::VecTestMeasure) = LebesgueBase()^2
 MeasureBase.insupport(::VecTestMeasure, x) = true
 MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s)
 
-# `@allocated` at top level reports a boxed result on Julia 1.10:
-_allocated(f::F, args::Vararg{Any,N}) where {F,N} = @allocated f(args...)
+include("testutils.jl")
 
 @testset "logdensities" begin
     @testset "scalar variates" begin
@@ -128,10 +127,10 @@ _allocated(f::F, args::Vararg{Any,N}) where {F,N} = @allocated f(args...)
         m3 = StdNormal()^static(3)
         xs = @SVector randn(3)
         @test @inferred(logdensityof(m3, xs)) ≈ sum(stdnormal_ld, xs)
-        @test _allocated(logdensityof, m3, xs) == 0
+        @test allocations_of(logdensityof, m3, xs) == 0
         xd = randn(3)
         @test @inferred(logdensityof(StdNormal()^3, xd)) ≈ sum(stdnormal_ld, xd)
-        @test _allocated(logdensityof, StdNormal()^3, xd) == 0
+        @test allocations_of(logdensityof, StdNormal()^3, xd) == 0
         Xs = @SMatrix randn(3, 4)
         @test @inferred(logdensities(m3, Xs)) ≈ vec(sum(stdnormal_ld.(Xs), dims = 1))
         @test logdensities(m3, Xs) isa SVector{4}
@@ -151,7 +150,7 @@ _allocated(f::F, args::Vararg{Any,N}) where {F,N} = @allocated f(args...)
         @test @inferred(logdensities(w, X)) ≈ [logdensityof(w, x) for x in eachcol(X)]
         xw = randn(3)
         @test @inferred(logdensityof(w, xw)) ≈ log(0.3) + sum(stdnormal_ld, xw)
-        @test _allocated(logdensityof, w, xw) == 0
+        @test allocations_of(logdensityof, w, xw) == 0
 
         ms = [weightedmeasure(log(i), StdNormal()) for i in 1:4]
         prod4 = productmeasure(ms)
@@ -159,7 +158,7 @@ _allocated(f::F, args::Vararg{Any,N}) where {F,N} = @allocated f(args...)
         @test @inferred(MeasureBase.mspace_elsize(prod4)) == (4,)
         xp = randn(4)
         @test @inferred(logdensityof(prod4, xp)) ≈ sum(log(i) + stdnormal_ld(xp[i]) for i in 1:4)
-        @test _allocated(logdensityof, prod4, xp) == 0
+        @test allocations_of(logdensityof, prod4, xp) == 0
         Xp = randn(4, 7)
         @test @inferred(logdensities(prod4, Xp)) ≈ [logdensityof(prod4, x) for x in eachcol(Xp)]
         @test logdensities(prod4, sliced(Xp, 1)) ≈ logdensities(prod4, Xp)
diff --git a/test/testutils.jl b/test/testutils.jl
new file mode 100644
index 00000000..5fa9c86f
--- /dev/null
+++ b/test/testutils.jl
@@ -0,0 +1,9 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+# Allocations of `f(args...)` after a warm-up call, measured inside a
+# function since `@allocated` at top level reports a boxed result on
+# Julia 1.10:
+function allocations_of(f::F, args::Vararg{Any,N}) where {F,N}
+    f(args...)
+    @allocated f(args...)
+end
diff --git a/test/transport.jl b/test/transport.jl
index ea3c8614..d2c31ae1 100644
--- a/test/transport.jl
+++ b/test/transport.jl
@@ -12,6 +12,8 @@ using LogExpFunctions: logit
 using ArraysOfArrays: sliced, flatview, fused
 using JLArrays
 
+include("testutils.jl")
+
 @testset "transport_to" begin
     for (f, μ) in [
         (logit, StdUniform())
@@ -90,11 +92,11 @@ using JLArrays
     @testset "scalar and static transports" begin
         f = transport_to(StdNormal(), StdUniform())
         @test @inferred(f(0.3)) isa Float64
-        @test @allocated(f(0.3)) == 0
+        @test allocations_of(f, 0.3) == 0
         g = transport_to(StdExponential()^static(3), StdNormal()^static(3))
         xs = SVector(0.1, -0.4, 2.0)
         @test @inferred(g(xs)) isa SVector{3,Float64}
-        @test @allocated(g(xs)) == 0
+        @test allocations_of(g, xs) == 0
         @test inverse(g)(g(xs)) ≈ xs
         h = transport_to(StdNormal()^3, StdUniform()^3)
         @test h(Float32[0.1, 0.5, 0.9]) isa Vector{Float32}

From f2c6d393f7cebf72f5165d0a31307975edd1768e Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 18 Sep 2026 19:19:26 +0200
Subject: [PATCH 104/122] Test transparency for fixed-size arrays

FixedSizeArrays becomes a test dependency: fixed-size inputs to
densities and transports give fixed-size outputs, since the kernels
allocate via similar.

Created by generative AI.
---
 redesign.md               |  5 +++-
 test/Project.toml         |  1 +
 test/fixed_size_arrays.jl | 48 +++++++++++++++++++++++++++++++++++++++
 test/runtests.jl          |  1 +
 4 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 test/fixed_size_arrays.jl

diff --git a/redesign.md b/redesign.md
index 518c32e0..7e571701 100644
--- a/redesign.md
+++ b/redesign.md
@@ -162,7 +162,10 @@ products gives plain arrays); vcat-combined and bind variates are flat;
 
 ## Verification
 
-Full suite (Aqua, extensions, doctests) on CPU with JLArrays cases.
+Full suite (Aqua, extensions, doctests) on CPU with JLArrays cases and
+a FixedSizeArrays transparency check (fixed-size inputs give fixed-size
+outputs; FixedSizeArrays stays a test dependency, allocating fixed-size
+variates by default is a HeterogeneousComputing decision for later).
 `test/test_reactant.jl` runs as part of the suite on 64-bit Linux and
 macOS with stable Julia, adding Reactant on demand as MGVI does (backend
 via `MEASUREBASE_REACTANT_BACKEND`). `test/cuda` is opt-in. Both run
diff --git a/test/Project.toml b/test/Project.toml
index e4ed8cc0..c518a35a 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -10,6 +10,7 @@ DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
diff --git a/test/fixed_size_arrays.jl b/test/fixed_size_arrays.jl
new file mode 100644
index 00000000..6006f24e
--- /dev/null
+++ b/test/fixed_size_arrays.jl
@@ -0,0 +1,48 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using MeasureBase
+using MeasureBase: StdNormal, StdUniform, StdExponential, logdensities, weightedmeasure, productmeasure, pushfwd
+using MeasureBase: batched_transport_to_std, batched_transport_from_std
+using InverseFunctions: inverse
+using FixedSizeArrays: FixedSizeArrayDefault
+using ArraysOfArrays: flatview, sliced
+using AffineMaps: Mul
+using Distributions: Normal
+
+# Fixed-size inputs give fixed-size outputs, the kernels allocate via
+# `similar` and never fall back to plain arrays:
+@testset "fixed-size arrays" begin
+    fixed(A) = FixedSizeArrayDefault(A)
+    isfixed(A) = A isa FixedSizeArrayDefault
+    isfixed(A::Union{SubArray,Base.ReshapedArray}) = isfixed(parent(A))
+    X = fixed(randn(3, 20))
+    x = fixed(randn(3))
+    m3 = StdNormal()^3
+
+    @testset "densities" begin
+        ℓ = logdensities(m3, X)
+        @test isfixed(ℓ) && ℓ ≈ logdensities(m3, Array(X))
+        @test logdensityof(m3, x) ≈ logdensityof(m3, Array(x))
+        @test isfixed(logdensities(m3, sliced(X, Val(1))))
+        @test isfixed(logdensities((StdNormal()^2)^3, fixed(randn(2, 3, 4))))
+        @test isfixed(logdensities(weightedmeasure(0.3, m3), X))
+        P = productmeasure(fixed([pushfwd(Mul(s), StdNormal()) for s in (1.0, 2.0, 3.0)]))
+        @test isfixed(logdensities(P, X))
+        Pn = productmeasure(fixed([Normal(μ, 1.0) for μ in (0.0, 1.0, 2.0)]))
+        @test isfixed(logdensities(Pn, X)) && logdensities(Pn, X) ≈ logdensities(Pn, Array(X))
+    end
+
+    @testset "transports" begin
+        f = transport_to(StdUniform()^3, m3)
+        y = f(x)
+        @test isfixed(y) && y ≈ f(Array(x))
+        Y = f.(X)
+        @test isfixed(flatview(Y)) && flatview(Y) ≈ flatview(f.(Array(X)))
+        @test flatview(inverse(f).(Y)) ≈ X
+        @test isfixed(batched_transport_to_std(StdNormal, StdExponential()^3, fixed(rand(3, 5))))
+        @test isfixed(batched_transport_from_std(StdNormal, StdExponential()^3, fixed(randn(3, 5))))
+        @test isfixed(MeasureBase.convert_realtype(Float32, x))
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 55b9d679..e75294d3 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -19,6 +19,7 @@ include("shape_contract.jl")
 include("logdensities.jl")
 include("structured_batches.jl")
 include("batched_regressions.jl")
+include("fixed_size_arrays.jl")
 include("numtype.jl")
 include("transport.jl")
 include("transport_batched.jl")

From d84cf94d95de8d56515f39996e593fc4ac258177 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sat, 19 Sep 2026 05:19:43 +0200
Subject: [PATCH 105/122] Adopt one convention for variates outside the support

Variates of the right shape never throw. Densities are -Inf outside the
support, including infinite values and non-integers for counting-based
measures, transports are NaN outside the support of the source measure,
and wrong shapes throw an ArgumentError from checked_arg. The standard
measure transports, Half and the wrapped distribution families mask
their results and use abs, clamp and min guards so that the formulas
stay total on the CPU and identical on devices. Infinite values lie
outside the support of wrapped univariate distributions. Documented in
logdensityof, transport_to and checked_arg, with tests through the
combinators on CPU and JLArrays.

Created by generative AI.
---
 .../distribution_measure.jl                   |  6 +-
 ext/MeasureBaseDistributionsExt/families.jl   | 23 +++----
 .../multivariate.jl                           | 20 +++++--
 ext/MeasureBaseDistributionsExt/univariate.jl | 14 ++++-
 redesign.md                                   | 28 +++++++--
 src/combinators/half.jl                       |  5 +-
 src/density-core.jl                           | 14 +++++
 src/getdof.jl                                 |  3 +
 src/standard/stdconvert.jl                    |  6 +-
 src/standard/stdexponential.jl                |  4 +-
 src/standard/stdlogistic.jl                   |  4 +-
 src/standard/stdnormal.jl                     |  4 +-
 src/standard/stduniform.jl                    |  5 +-
 src/transport.jl                              |  8 +++
 test/distributions/test_device_kernels.jl     | 33 +++++++++-
 test/runtests.jl                              |  1 +
 test/support_conventions.jl                   | 60 +++++++++++++++++++
 17 files changed, 198 insertions(+), 40 deletions(-)
 create mode 100644 test/support_conventions.jl

diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index f0c152cb..ddd6aced 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -68,7 +68,11 @@ for (bhead, phead) in ((:batched_logdensityof_impl, :logdensityof_impl), (:batch
     end
 end
 @inline MeasureBase.unsafe_logdensityof(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x)
-@inline MeasureBase.insupport(m::DistributionMeasure, x) = Distributions.insupport(m.obj, x)
+@inline MeasureBase.insupport(m::DistributionMeasure, x) = Distributions.insupport(m.obj, x) & _finite_variate(m.obj, x)
+# Infinite values lie outside the support of univariate distributions,
+# where Distributions may evaluate to NaN:
+@inline _finite_variate(::Distribution{Univariate}, x) = isfinite(x)
+@inline _finite_variate(::Distribution, x) = true
 
 @inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate{0},<:Continuous}) = Lebesgue()
 @inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate,<:Continuous}) = Lebesgue()^size(m.obj)
diff --git a/ext/MeasureBaseDistributionsExt/families.jl b/ext/MeasureBaseDistributionsExt/families.jl
index a0f757e0..ffd24a9c 100644
--- a/ext/MeasureBaseDistributionsExt/families.jl
+++ b/ext/MeasureBaseDistributionsExt/families.jl
@@ -13,7 +13,7 @@ const _Families = Union{Normal,Uniform,Exponential,Logistic,Cauchy,Laplace,LogNo
 # not throw outside of the support, where their results are masked:
 @inline MeasureBase.logdensity_def(m::AsMeasure{<:_Families}, x) = _family_logd(m.obj, x)
 @inline MeasureBase.unsafe_logdensityof(m::AsMeasure{<:_Families}, x) = _family_logd(m.obj, x)
-@inline MeasureBase.insupport(m::AsMeasure{<:_Families}, x) = _family_insupport(m.obj, x)
+@inline MeasureBase.insupport(m::AsMeasure{<:_Families}, x) = _family_insupport(m.obj, x) & _finite_variate(m.obj, x)
 
 # `c * log(y)`, zero for `c == 0` also where `y == 0`:
 @inline _clog(c, y) = ifelse(iszero(c), zero(c * log(one(y))), c * log(y))
@@ -104,7 +104,8 @@ end
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Cauchy}) = StdUniform
 @inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Cauchy, x) = 1 // 2 + atan((x - d.μ) / d.σ) / π
-@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Cauchy, p) = muladd(d.σ, tan(π * (p - 1 // 2)), d.μ)
+@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Cauchy, p) =
+    _nan_outside(StdUniform, p, muladd(d.σ, tan(π * (p - 1 // 2)), d.μ))
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Laplace}) = StdUniform
 @inline function MeasureBase.transport_to_std(::Type{StdUniform}, d::Laplace, x)
@@ -113,21 +114,23 @@ end
 end
 @inline function MeasureBase.transport_from_std(::Type{StdUniform}, d::Laplace, p)
     u = p - 1 // 2
-    muladd(-d.θ * sign(u), log1p(-2 * abs(u)), d.μ)
+    _nan_outside(StdUniform, p, muladd(-d.θ * sign(u), log1p(-min(2 * abs(u), one(u))), d.μ))
 end
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:LogNormal}) = StdNormal
-@inline MeasureBase.transport_to_std(::Type{StdNormal}, d::LogNormal, x) = (log(x) - d.μ) / d.σ
+@inline MeasureBase.transport_to_std(::Type{StdNormal}, d::LogNormal, x) = _nan_outside(d, x, (log(abs(x)) - d.μ) / d.σ)
 @inline MeasureBase.transport_from_std(::Type{StdNormal}, d::LogNormal, z) = exp(muladd(d.σ, z, d.μ))
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Weibull}) = StdExponential
-@inline MeasureBase.transport_to_std(::Type{StdExponential}, d::Weibull, x) = (x / d.θ)^d.α
-@inline MeasureBase.transport_from_std(::Type{StdExponential}, d::Weibull, z) = d.θ * z^(1 / d.α)
+@inline MeasureBase.transport_to_std(::Type{StdExponential}, d::Weibull, x) = _nan_outside(d, x, abs(x / d.θ)^d.α)
+@inline MeasureBase.transport_from_std(::Type{StdExponential}, d::Weibull, z) = _nan_outside(StdExponential, z, d.θ * abs(z)^(1 / d.α))
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Gamma}) = StdUniform
-@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Gamma, x) = _gamma_cdf(d.α, x / d.θ)
-@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Gamma, p) = d.θ * _gamma_quantile(d.α, p)
+@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Gamma, x) = _nan_outside(d, x, _gamma_cdf(d.α, abs(x / d.θ)))
+@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Gamma, p) = _nan_outside(StdUniform, p, d.θ * _gamma_quantile(d.α, _unit_clamp(p)))
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Beta}) = StdUniform
-@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Beta, x) = _beta_cdf(d.α, d.β, x)
-@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Beta, p) = _beta_quantile(d.α, d.β, p)
+@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Beta, x) = _nan_outside(d, x, _beta_cdf(d.α, d.β, _unit_clamp(x)))
+@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Beta, p) = _nan_outside(StdUniform, p, _beta_quantile(d.α, d.β, _unit_clamp(p)))
+
+@inline _unit_clamp(x) = clamp(x, zero(x), one(x))
diff --git a/ext/MeasureBaseDistributionsExt/multivariate.jl b/ext/MeasureBaseDistributionsExt/multivariate.jl
index 6ca640f9..b4262882 100644
--- a/ext/MeasureBaseDistributionsExt/multivariate.jl
+++ b/ext/MeasureBaseDistributionsExt/multivariate.jl
@@ -22,6 +22,9 @@
 @inline _rows(Z::AbstractMatrix, r) = view(Z, r, :)
 @inline _masked(ℓ::Number, ins) = ifelse(ins, ℓ, oftype(ℓ, -Inf))
 @inline _masked(ℓ::AbstractArray, ins) = ifelse.(ins, ℓ, eltype(ℓ)(-Inf))
+@inline _nan_columns(Z::AbstractArray, ins) = ifelse.(_as_row(ins), Z, eltype(Z)(NaN))
+@inline _as_row(ins::Bool) = ins
+@inline _as_row(ins::AbstractVector) = reshape(ins, 1, :)
 
 
 # Multivariate normal: densities via the Cholesky factor of the covariance.
@@ -75,14 +78,17 @@ for bhead in (:batched_logdensityof_impl, :batched_logdensity_def)
         d = m.obj
         Xc = _as_columns(X)
         ℓ = _column_sums(identity, _clog.(d.alpha .- 1, abs.(Xc))) .- d.lmnB
-        tol = sqrt(eps(float(eltype(X))))
-        ins = _column_all(Xc .>= 0) .& (abs.(_column_sums(identity, Xc) .- 1) .<= tol)
-        _batch_results(_masked(ℓ, ins), X)
+        _batch_results(_masked(ℓ, _simplex_mask(Xc)), X)
     end
 end
 MeasureBase.logdensity_def(m::DirichletMeasure, x::AbstractVector) = MeasureBase.batched_logdensity_def(m, x)
 MeasureBase.unsafe_logdensityof(m::DirichletMeasure, x::AbstractVector) = MeasureBase.batched_logdensityof_impl(m, x)
 
+@inline function _simplex_mask(Xc::AbstractArray)
+    tol = sqrt(eps(float(eltype(Xc))))
+    _column_all(Xc .>= 0) .& (abs.(_column_sums(identity, Xc) .- 1) .<= tol)
+end
+
 # The stick-breaking Beta parameters, for the first `K - 1` components:
 @inline _stick_breaking_params(d::Dirichlet) = (_dropfront(_rev_cumsum(d.alpha)), _dropback(d.alpha))
 
@@ -94,17 +100,19 @@ function MeasureBase.batched_transport_to_std(::Type{StdUniform}, d::Dirichlet,
     # The remaining mass before each component is the mass after it plus
     # the component itself:
     beta_v = _rows(rem, 1:(K - 1)) ./ (_rows(rem, 1:(K - 1)) .+ _rows(Xc, 1:(K - 1)))
-    _from_columns(_beta_cdf.(αs, βs, beta_v), X)
+    Z = _beta_cdf.(αs, βs, _unit_clamp.(beta_v))
+    _from_columns(_nan_columns(Z, _simplex_mask(Xc)), X)
 end
 
 function MeasureBase.batched_transport_from_std(::Type{StdUniform}, d::Dirichlet, Z::AbstractArray)
     K = length(d)
     αs, βs = _stick_breaking_params(d)
-    beta_v = _beta_quantile.(αs, βs, _as_columns(Z))
+    Zc = _as_columns(Z)
+    beta_v = _beta_quantile.(αs, βs, _unit_clamp.(Zc))
     cp = cumprod(beta_v; dims = 1)
     # Each component takes what its Beta variate leaves of the remaining mass:
     X = vcat(1 .- _rows(cp, 1:1), _rows(cp, 1:(K - 2)) .- _rows(cp, 2:(K - 1)), _rows(cp, (K - 1):(K - 1)))
-    _from_columns(X, Z)
+    _from_columns(_nan_columns(X, _column_all((Zc .>= 0) .& (Zc .<= 1))), Z)
 end
 MeasureBase.transport_to_std(::Type{StdUniform}, d::Dirichlet, x) = MeasureBase.batched_transport_to_std(StdUniform, d, x)
 MeasureBase.transport_from_std(::Type{StdUniform}, d::Dirichlet, z) = MeasureBase.batched_transport_from_std(StdUniform, d, z)
diff --git a/ext/MeasureBaseDistributionsExt/univariate.jl b/ext/MeasureBaseDistributionsExt/univariate.jl
index 8d4b26da..1f318379 100644
--- a/ext/MeasureBaseDistributionsExt/univariate.jl
+++ b/ext/MeasureBaseDistributionsExt/univariate.jl
@@ -94,8 +94,12 @@ end
     convert(_result_numtype(d, z), x)
 end
 
+# Transports outside the support of the source give NaN, the formulas
+# must not throw outside the support (see `MeasureBase._nan_outside`):
+@inline _nan_outside(d::Distribution, x, y) = MeasureBase._nan_outside(asmeasure(d), x, y)
+@inline _nan_outside(::Type{S}, z, y) where {S<:StdMeasure} = MeasureBase._nan_outside(S(), z, y)
+
 for (D, S) in [
-    (Uniform, StdUniform),
     (Logistic, StdLogistic),
     (Normal, StdNormal)
 ]
@@ -106,11 +110,15 @@ for (D, S) in [
     end
 end
 
+@inline MeasureBase.preferred_stdmeasure(::Type{<:Uniform}) = StdUniform
+@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Uniform, x) = _nan_outside(d, x, _affine_to_std(d, x))
+@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Uniform, z) = _nan_outside(StdUniform, z, _std_to_affine(d, z))
+
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Exponential}) = StdExponential
 @inline MeasureBase.transport_to_std(::Type{StdExponential}, d::Exponential, x) =
-    convert(_result_numtype(d, x), Distributions.scale(d) \ x)
+    _nan_outside(d, x, convert(_result_numtype(d, x), Distributions.scale(d) \ x))
 @inline MeasureBase.transport_from_std(::Type{StdExponential}, d::Exponential, z) =
-    convert(_result_numtype(d, z), Distributions.scale(d) * z)
+    _nan_outside(StdExponential, z, convert(_result_numtype(d, z), Distributions.scale(d) * z))
 
 
 # Affine transformed distributions transport via the underlying distribution:
diff --git a/redesign.md b/redesign.md
index 7e571701..be52602e 100644
--- a/redesign.md
+++ b/redesign.md
@@ -184,10 +184,30 @@ locally on the GB10, green at HEAD except one expected-broken CUDA case
   HeterogeneousComputing has no Reactant compute unit; JLArrays has no
   RNG and no triangular solves; Reactant rejects traced `VectorOfArrays`
   and empty batches.
-- Decisions pending: one convention for out-of-support inputs (NaN mask
-  vs. DomainError vs. AssertionError), `Half` tails via log-ccdf, device
-  random variate infrastructure and `rand!`, Tier-1 static variates,
-  the `smart-constructors.jl` review (location-scale arrays as affine
+- Out-of-support convention (decided 2026-09-19): structural errors (wrong
+  rank, size, container or element kind) throw an `ArgumentError` from
+  `checked_arg` at the entry points; variates of the right shape never
+  throw: densities are `-Inf` outside the support (including non-integers
+  for counting-based measures and infinite values), transports are `NaN`
+  outside the support of the source (boundaries may map to `±Inf`), `NaN`
+  inputs give `NaN` or `-Inf`, relative densities keep `+Inf`/`-Inf`/`NaN`.
+  Kernels must not throw outside the support, since the masks evaluate
+  both branches (`abs`, `clamp`, `min` guards instead of `NaNMath`, which
+  isn't device-compatible). Downstream checks such as BAT's
+  `checked_logdensityof` stay downstream.
+- BAT's boundary tweaks (to discuss): adopted are infinite variates
+  outside the support of continuous wrapped distributions, `-Inf + Inf`
+  giving `-Inf` in pushforward densities and a zero Jacobian term where
+  both densities vanish. Not adopted: clamping uniform-direction inputs to
+  `[eps, 1 - eps]` (here `NaN` outside, `±Inf` or support edges at 0 and
+  1), replacing finite densities with infinite Jacobian terms by `-1e38`,
+  and re-evaluating densities an `eps` inside the support where
+  Distributions returns `NaN` (the family kernels are exact there).
+  Quantile results within `4 eps` of the support edges snap to the edges
+  only on the generic logistic path of wrapped distributions.
+- Decisions pending: `Half` tails via log-ccdf, device random variate
+  infrastructure and `rand!`, Tier-1 static variates, the
+  `smart-constructors.jl` review (location-scale arrays as affine
   pushforwards of powers), `_static_ndims` type-first vs. instance-first.
 - Polish before merge: docs pass, NEWS, history curation, version bump,
   remove this file.
diff --git a/src/combinators/half.jl b/src/combinators/half.jl
index c6ed31da..f793a947 100644
--- a/src/combinators/half.jl
+++ b/src/combinators/half.jl
@@ -46,9 +46,8 @@ function smf(μ::Half, x)
 end
 
 function invsmf(μ::Half, p)
-    @assert zero(p) ≤ p ≤ one(p)
-    invsmf(μ.parent, (p + 1) / 2)
+    _nan_outside(StdUniform(), p, invsmf(μ.parent, (min(p, one(p)) + 1) / 2))
 end
 
-@inline transport_to_std(::Type{StdUniform}, μ::Half, x) = smf(μ, x)
+@inline transport_to_std(::Type{StdUniform}, μ::Half, x) = _nan_outside(μ, x, smf(μ, x))
 @inline transport_from_std(::Type{StdUniform}, μ::Half, p) = invsmf(μ, p)
diff --git a/src/density-core.jl b/src/density-core.jl
index 71ef2752..e39f9809 100644
--- a/src/density-core.jl
+++ b/src/density-core.jl
@@ -27,6 +27,15 @@ To compute log-density relative to `basemeasure(m)` or *define* a log-density
 
 To compute a log-density relative to a specific base-measure, see
 `logdensity_rel`.
+
+# Extended help
+
+Variates of the right shape and element type never throw: outside the
+support of `m` the result is `-Inf`, also for non-integer values of
+measures over counting measures and for infinite values. `NaN` inputs give
+`NaN` or `-Inf`. Variates of the wrong shape throw an `ArgumentError`.
+Implementations of `logdensityof_impl` and `unsafe_logdensityof` must not
+throw outside the support, since support masks evaluate both branches.
 """
 @inline logdensityof(μ::AbstractMeasure, x) = _point_ld(logdensityof_impl, μ, x)
 
@@ -73,6 +82,11 @@ end
 
 @inline _checksupport(cond, result) = ifelse(_insupport_mask(cond), result, oftype(result, -Inf))
 
+# Transports of variates outside the support of the source measure give
+# NaN. Both branches are evaluated, formulas must not throw outside the
+# support:
+@inline _nan_outside(μ, x, y) = ifelse(_insupport_mask(insupport(μ, x)), y, oftype(y, NaN))
+
 """
     MeasureBase.logdensityof_with_rest(μ::AbstractMeasure, x)
 
diff --git a/src/getdof.jl b/src/getdof.jl
index aca22e67..fd9b4a20 100644
--- a/src/getdof.jl
+++ b/src/getdof.jl
@@ -148,6 +148,9 @@ struct NoArgCheck{MU,T} end
 
 Return `x` if `x` is a valid variate of `μ`, throw an `ArgumentError` if not,
 return `NoArgCheck{MU,T}()` if not check can be performed.
+
+Only the shape and type of `x` are checked, never its value: values
+outside the support are valid arguments of densities and transports.
 """
 function checked_arg end
 
diff --git a/src/standard/stdconvert.jl b/src/standard/stdconvert.jl
index 06fdd9c7..3868684b 100644
--- a/src/standard/stdconvert.jl
+++ b/src/standard/stdconvert.jl
@@ -13,8 +13,8 @@
     ifelse(z < zero(z), -log1p(-Φ(z)), -log(_normccdf(z)))
 end
 
-@inline function transport_def(::StdNormal, ::StdExponential, x)
-    ifelse(x < oftype(x, logtwo), Φinv(-expm1(-x)), -Φinv(exp(-x)))
+@inline function transport_def(::StdNormal, μ::StdExponential, x)
+    _nan_outside(μ, x, ifelse(x < oftype(x, logtwo), Φinv(-expm1(-x)), -Φinv(exp(-x))))
 end
 
 @inline transport_def(::StdLogistic, ::StdNormal, z) = _normlogcdf(z) - _normlogccdf(z)
@@ -23,7 +23,7 @@ end
     ifelse(l < zero(l), Φinv(logistic(l)), -Φinv(logistic(-l)))
 end
 
-@inline transport_def(::StdLogistic, ::StdExponential, x) = log(-expm1(-x)) + x
+@inline transport_def(::StdLogistic, μ::StdExponential, x) = _nan_outside(μ, x, log(abs(expm1(-x))) + x)
 
 @inline transport_def(::StdExponential, ::StdLogistic, l) = log1pexp(l)
 
diff --git a/src/standard/stdexponential.jl b/src/standard/stdexponential.jl
index 39043e4a..1bb89c24 100644
--- a/src/standard/stdexponential.jl
+++ b/src/standard/stdexponential.jl
@@ -18,8 +18,8 @@ end
 @inline logdensity_def(::StdExponential, x) = -x
 @inline basemeasure(::StdExponential) = LebesgueBase()
 
-@inline transport_def(::StdUniform, μ::StdExponential, x) = -expm1(-x)
-@inline transport_def(::StdExponential, μ::StdUniform, x) = -log1p(-x)
+@inline transport_def(::StdUniform, μ::StdExponential, x) = _nan_outside(μ, x, -expm1(-x))
+@inline transport_def(::StdExponential, μ::StdUniform, x) = _nan_outside(μ, x, -log1p(-min(x, one(x))))
 
 @inline rand_impl(ctx::GenContext, ::StdExponential) = randexp(get_rng(ctx), get_precision(ctx))
 @inline batched_rand_impl(ctx::GenContext, ::StdExponential, sz::Dims) = _randexp_bulk(ctx, sz)
diff --git a/src/standard/stdlogistic.jl b/src/standard/stdlogistic.jl
index 44b8a0b8..d1cfcd44 100644
--- a/src/standard/stdlogistic.jl
+++ b/src/standard/stdlogistic.jl
@@ -16,7 +16,9 @@ export StdLogistic
 @inline basemeasure(::StdLogistic) = LebesgueBase()
 
 @inline transport_def(::StdUniform, μ::StdLogistic, x) = logistic(x)
-@inline transport_def(::StdLogistic, μ::StdUniform, p) = logit(p)
+@inline transport_def(::StdLogistic, μ::StdUniform, p) = _nan_outside(μ, p, _logit_nan(p))
+# Total on the CPU, the mask supplies the NaN outside of [0, 1]:
+@inline _logit_nan(p) = log(abs(p) / abs(one(p) - p))
 
 @inline rand_impl(ctx::GenContext, ::StdLogistic) = logit(rand(get_rng(ctx), get_precision(ctx)))
 @inline batched_rand_impl(ctx::GenContext, ::StdLogistic, sz::Dims) = logit.(_rand_bulk(ctx, sz))
diff --git a/src/standard/stdnormal.jl b/src/standard/stdnormal.jl
index 989e645e..21312653 100644
--- a/src/standard/stdnormal.jl
+++ b/src/standard/stdnormal.jl
@@ -24,7 +24,7 @@ export StdNormal
 @inline batched_rand_impl(ctx::GenContext, ::StdNormal, sz::Dims) = _randn_bulk(ctx, sz)
 
 Φ(z) = erfc(-z * invsqrt2) / 2
-Φinv(p) = -erfcinv(2 * p) * sqrt2
+Φinv(p) = -erfcinv(2 * clamp(p, zero(p), one(p))) * sqrt2
 
 InverseFunctions.inverse(::typeof(Φ)) = Φinv
 InverseFunctions.inverse(::typeof(Φinv)) = Φ
@@ -35,5 +35,5 @@ invsmf(::StdNormal, p) = Φinv(p)
 smf(::StdNormal) = Φ
 invsmf(::StdNormal) = Φinv
 
-transport_def(::StdNormal, ::StdUniform, p) = Φinv(p)
+transport_def(::StdNormal, μ::StdUniform, p) = _nan_outside(μ, p, Φinv(p))
 transport_def(::StdUniform, ::StdNormal, x) = Φ(x)
diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl
index b5ec1afa..3c45f078 100644
--- a/src/standard/stduniform.jl
+++ b/src/standard/stduniform.jl
@@ -25,7 +25,4 @@ massof(::StdUniform, s::Interval) = massof(Lebesgue(0.0 .. 1.0), s)
 
 smf(::StdUniform, x) = clamp(x, zero(x), one(x))
 
-function invsmf(::StdUniform, p)
-    @assert zero(p) ≤ p ≤ one(p)
-    p
-end
+invsmf(d::StdUniform, p) = _nan_outside(d, p, p)
diff --git a/src/transport.jl b/src/transport.jl
index 9fde5ff0..ba8387be 100644
--- a/src/transport.jl
+++ b/src/transport.jl
@@ -52,6 +52,14 @@ export transport_to
 
 Transport `x` from the measure `μ` to the measure `ν`, equivalent to
 `transport_to(ν, μ)(x)`.
+
+# Extended help
+
+Variates of the right shape never throw: outside the support of `μ` the
+result is `NaN` (elementwise for powers and products), values at the
+boundary of the support may map to infinite values. Variates of the wrong
+shape throw an `ArgumentError`. Transport implementations must not throw
+outside the support, since the `NaN` masks evaluate both branches.
 """
 transport_to(ν, μ, x) = transport_to(ν, μ)(x)
 
diff --git a/test/distributions/test_device_kernels.jl b/test/distributions/test_device_kernels.jl
index c4b53932..c75dffcf 100644
--- a/test/distributions/test_device_kernels.jl
+++ b/test/distributions/test_device_kernels.jl
@@ -32,7 +32,7 @@ using JLArrays
         @test all(map((a, b) -> a == b || a ≈ b || (isnan(a) && isnan(b)), logdensityof.(Ref(m), xs), ℓ_ref))
         @test logdensities(m, xs) ≈ ℓ_ref nans = true
         @test Array(logdensities(m, JLArray(xs))) ≈ ℓ_ref nans = true
-        @test insupport.(Ref(m), xs) == Distributions.insupport.(d, xs)
+        @test insupport.(Ref(m), xs) == (Distributions.insupport.(d, xs) .& isfinite.(xs))
 
         if d isa ContinuousUnivariateDistribution
             x = rand(stblrng(), d, 12)
@@ -112,3 +112,34 @@ using JLArrays
         @test size(Xr) == (4, 200) && all(sum(Xr; dims = 1) .≈ 1)
     end
 end
+
+@testset "wrapped distributions outside the support" begin
+    for d in (Uniform(-1.0, 2.5), Exponential(0.7), LogNormal(0.2, 0.6), Weibull(1.4, 0.9), Gamma(2.3, 1.2), Beta(2.5, 3.5))
+        m = asmeasure(d)
+        x = minimum(d) - 1
+        @test logdensityof(m, x) == -Inf
+        @test isnan(transport_to(StdNormal(), m)(x))
+        @test isnan(transport_to(StdNormal(), m)(Inf)) || d isa Union{Exponential,LogNormal,Weibull,Gamma}
+        Y = Array(transport_to(StdNormal(), m).(JLArray([x, mean(d)])))
+        @test isnan(Y[1]) && !isnan(Y[2])
+    end
+    for d in (Cauchy(0.1, 0.8), Laplace(-0.4, 1.1), Gamma(2.3, 1.2), Beta(2.5, 3.5), Uniform(-1.0, 2.5))
+        f = transport_to(asmeasure(d), StdUniform())
+        @test isnan(f(1.5)) && isnan(f(-0.5)) && !isnan(f(0.3))
+        @test isequal(Array(f.(JLArray([1.5, 0.3]))), f.([1.5, 0.3]))
+    end
+    for d in (Exponential(0.7), Weibull(1.4, 0.9))
+        @test isnan(transport_to(asmeasure(d), StdExponential())(-1.0))
+    end
+    @test logdensityof(asmeasure(Rayleigh(2.0)), Inf) == -Inf
+    @test logdensityof(asmeasure(Poisson(2.7)), Inf) == -Inf
+    @test logdensityof(asmeasure(Poisson(2.7)), 1.5) == -Inf
+    @test logdensityof(asmeasure(Poisson(2.7)), -1.0) == -Inf
+    @test logdensityof(asmeasure(Bernoulli(0.3)), 0.5) == -Inf
+    md = asmeasure(Dirichlet([2.0, 3.0, 4.0, 1.5]))
+    @test all(isnan, transport_to(StdUniform()^3, md)([0.5, 0.5, 0.5, 0.5]))
+    @test all(isnan, transport_to(md, StdUniform()^3)([1.5, 0.5, 0.5]))
+    @test logdensityof(md, [0.5, 0.5, 0.5, 0.5]) == -Inf
+    Zd = flatview(transport_to(StdUniform()^3, md).(sliced([0.5 0.1; 0.5 0.2; 0.5 0.3; 0.5 0.4], Val(1))))
+    @test all(isnan, Zd[:, 1]) && !any(isnan, Zd[:, 2])
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index e75294d3..4fdbc12d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -20,6 +20,7 @@ include("logdensities.jl")
 include("structured_batches.jl")
 include("batched_regressions.jl")
 include("fixed_size_arrays.jl")
+include("support_conventions.jl")
 include("numtype.jl")
 include("transport.jl")
 include("transport_batched.jl")
diff --git a/test/support_conventions.jl b/test/support_conventions.jl
new file mode 100644
index 00000000..9cea29c4
--- /dev/null
+++ b/test/support_conventions.jl
@@ -0,0 +1,60 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+using Test
+
+using MeasureBase
+using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic, Half, logdensities
+using MeasureBase: weightedmeasure, productmeasure, mcombine
+using JLArrays: JLArray
+
+# Variates of the right shape never throw: densities are -Inf and
+# transports are NaN outside the support, wrong shapes throw.
+@testset "support conventions" begin
+    outside = Dict(StdUniform() => -0.5, StdExponential() => -1.0, Half(StdNormal()) => -1.0)
+
+    @testset "densities" begin
+        for (μ, x) in outside
+            @test logdensityof(μ, x) == -Inf
+            @test logdensityof(weightedmeasure(0.3, μ), x) == -Inf
+            X = [x 0.5 0.5; 0.5 x 0.5; 0.5 0.5 0.5]
+            ℓ = logdensities(μ^3, X)
+            @test ℓ[1:2] == [-Inf, -Inf] && isfinite(ℓ[3])
+            @test Array(logdensities(μ^3, JLArray(X))) == ℓ
+            @test Array(logdensities(μ, JLArray(vec(X)))) == logdensities(μ, vec(X))
+        end
+        @test logdensityof(StdUniform(), Inf) == -Inf
+        @test logdensityof(StdExponential(), Inf) == -Inf
+        @test logdensityof(StdNormal(), -Inf) == -Inf
+        @test logdensityof(productmeasure((StdUniform(), StdExponential())), (0.5, -1.0)) == -Inf
+        @test logdensityof(mcombine(vcat, StdUniform()^2, StdExponential()^1), [0.5, 1.5, 0.5]) == -Inf
+    end
+
+    @testset "transports" begin
+        stds = (StdUniform(), StdExponential(), StdLogistic(), StdNormal())
+        for (μ, x) in outside, ν in stds
+            μ === ν && continue
+            f = transport_to(ν, μ)
+            @test isnan(f(x))
+            @test isnan(transport_to(ν^2, μ^2)([x, 0.5])[1])
+            Y = f.([x, 0.5, 0.5])
+            @test isnan(Y[1]) && !isnan(Y[2])
+            @test isequal(Array(f.(JLArray([x, 0.5, 0.5]))), Y)
+            @test isnan(transport_to(μ, ν)(NaN))
+        end
+        for ν in (StdExponential(), StdLogistic(), StdNormal(), Half(StdNormal()))
+            @test isnan(transport_to(ν, StdUniform())(1.5))
+            @test isnan(transport_to(ν, StdUniform())(-0.5))
+        end
+        @test transport_to(StdExponential(), StdUniform())(1.0) == Inf
+        @test transport_to(StdNormal(), StdUniform())(0.0) == -Inf
+        @test !isnan(transport_to(StdUniform(), StdNormal())(-37.0))
+        @test !isnan(transport_to(StdUniform(), StdLogistic())(-800.0))
+    end
+
+    @testset "wrong shapes throw" begin
+        @test_throws ArgumentError logdensityof(StdNormal()^3, randn(2))
+        @test_throws ArgumentError logdensities(StdNormal()^3, randn(2, 5))
+        @test_throws ArgumentError transport_to(StdUniform()^3, StdNormal()^3)(randn(2))
+        @test_throws ArgumentError logdensityof(StdNormal(), randn(2))
+    end
+end

From d89e3e987cb8f2acd8bcb7a6498a164f329b14be Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sat, 19 Sep 2026 05:19:43 +0200
Subject: [PATCH 106/122] Bump version to 0.15.0

Created by generative AI.
---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 5d59432f..db29f9d7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 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]

From 4269efa066d21bda815943b46bb26323f3da129d Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sat, 19 Sep 2026 06:06:21 +0200
Subject: [PATCH 107/122] Support PropertyFunctions 0.3

The marginalization map is built with the PropSelFunction constructor,
whose type parameters changed in PropertyFunctions 0.3.

Created by generative AI.
---
 Project.toml                        | 2 +-
 src/combinators/implicitlymapped.jl | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index db29f9d7..a8e8b924 100644
--- a/Project.toml
+++ b/Project.toml
@@ -93,7 +93,7 @@ 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"
diff --git a/src/combinators/implicitlymapped.jl b/src/combinators/implicitlymapped.jl
index 36a72a93..721d7268 100644
--- a/src/combinators/implicitlymapped.jl
+++ b/src/combinators/implicitlymapped.jl
@@ -227,7 +227,7 @@ export Marginalized
 implicit_origin(mapped::Marginalized) = mapped.obj
 
 function explicit_mapfunc(::Marginalized, obs::NamedTuple{names}) where {names}
-    PropSelFunction{names,names}()
+    PropSelFunction(names...)
 end
 function pushfwd(f::PropSelFunction, mu::ProductMeasure{<:NamedTuple}, ::PushfwdRootMeasure)
     productmeasure(f(marginals(mu)))

From 0d4698b3aeef89e9467cda1173025edd442a855d Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sat, 19 Sep 2026 13:07:44 +0200
Subject: [PATCH 108/122] Count every component base measure in superposition
 densities

The base measure of a superposition sums the base measures of all
components unconditionally, so the Radon-Nikodym derivative of a
component with respect to it must count every base measure with mass at
the point, not only those of components whose support contains it.
Otherwise densities came out too large where a component vanishes, e.g.
by log(2) for a superposition of a normal and a uniform measure outside
the unit interval.

Created by generative AI.
---
 src/combinators/superpose.jl  |  4 ++--
 test/combinators/superpose.jl | 10 ++++++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index 9042d283..db7a8950 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -84,8 +84,8 @@ function logdensity_def(s::SuperpositionMeasure, x)
     αs = map(basemeasure, cs)
     terms = map(cs, αs) do cᵢ, αᵢ
         ℓᵢ = _dynamic_logd(logdensity_def(cᵢ, x), x)
-        log_dΣα_dαᵢ = _logsumexp_components(map(cs, αs) do cⱼ, αⱼ
-            _masked_logd(logdensity_rel(αⱼ, αᵢ, x), insupport(cⱼ, x))
+        log_dΣα_dαᵢ = _logsumexp_components(map(αs) do αⱼ
+            _masked_logd(logdensity_rel(αⱼ, αᵢ, x), insupport(αⱼ, x))
         end)
         _masked_logd(ℓᵢ - log_dΣα_dαᵢ, insupport(cᵢ, x))
     end
diff --git a/test/combinators/superpose.jl b/test/combinators/superpose.jl
index 753d5b8c..2369ae5e 100644
--- a/test/combinators/superpose.jl
+++ b/test/combinators/superpose.jl
@@ -9,13 +9,19 @@ using MeasureBase: superpose, weightedmeasure, StdNormal
     μs = μ + ν
     @test μs isa SuperpositionMeasure{<:Tuple{Dirac,Dirac}}
     @test μs == SuperpositionMeasure((μ, ν)) == superpose(μ, ν)
-    @test density_def(μs, 0) == 1.0
+    @test density_def(μs, 0) == 0.5
     @test basemeasure(μs) == CountingBase() + CountingBase()
+    @test densityof(μs, 0) == 1.0
 
     μs = SuperpositionMeasure([μ, ν])
     @test μs isa SuperpositionMeasure{<:AbstractVector{<:AbstractMeasure}}
-    @test density_def(μs, 0) == 1.0
+    @test density_def(μs, 0) == 0.5
     @test basemeasure(μs) == weightedmeasure(log(2), CountingBase())
+    @test densityof(μs, 0) == 1.0
+    # Base measures of components count wherever they have mass, not only
+    # where the component itself does:
+    @test logdensityof(superpose(StdNormal(), StdUniform()), -1.0) ≈ logdensityof(StdNormal(), -1.0)
+    @test logdensityof(superpose(StdNormal(), StdUniform()), 0.5) ≈ log(exp(logdensityof(StdNormal(), 0.5)) + 1)
 
     # Dirac equality is not decidable from types, so no weighted collapse:
     μ2 = μ + μ

From 10989a6ae68d1baff53fcb5ff2842b44d825a988 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sat, 19 Sep 2026 13:07:44 +0200
Subject: [PATCH 109/122] Hash measures and transports by value

AsMeasure, Dirac, product measures and transport functions compare by
value, so they hash by value as well, which keeps equal measures with
array-valued parameters usable as dictionary keys and cache witnesses.

Created by generative AI.
---
 src/MeasureBase.jl         | 1 +
 src/combinators/product.jl | 1 +
 src/primitives/dirac.jl    | 1 +
 src/transport.jl           | 1 +
 test/test_basics.jl        | 8 ++++++++
 5 files changed, 12 insertions(+)

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index a0d29857..8d5dafb9 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -131,6 +131,7 @@ ConstructionBase.constructorof(::Type{<:AsMeasure}) = _asmeasure
 _asmeasure(obj) = AsMeasure{typeof(obj)}(obj)
 
 Base.:(==)(a::AsMeasure, b::AsMeasure) = a.obj == b.obj
+Base.hash(a::AsMeasure, h::UInt) = hash(a.obj, hash(:AsMeasure, h))
 Base.isapprox(a::AsMeasure, b::AsMeasure; kwargs...) = isapprox(a.obj, b.obj; kwargs...)
 
 function Pretty.quoteof(d::M) where {M<:AbstractMeasure}
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 3669f358..e370e62a 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -23,6 +23,7 @@ export marginals
 function Base.:(==)(a::AbstractProductMeasure, b::AbstractProductMeasure)
     marginals(a) == marginals(b)
 end
+Base.hash(a::AbstractProductMeasure, h::UInt) = hash(marginals(a), hash(:AbstractProductMeasure, h))
 Base.length(μ::AbstractProductMeasure) = length(marginals(μ))
 Base.size(μ::AbstractProductMeasure) = size(marginals(μ))
 
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 08a27a83..f5025bb9 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -10,6 +10,7 @@ function Pretty.tile(d::Dirac)
 end
 
 Base.:(==)(a::Dirac, b::Dirac) = a.x == b.x
+Base.hash(a::Dirac, h::UInt) = hash(a.x, hash(:Dirac, h))
 Base.isapprox(a::Dirac, b::Dirac; kwargs...) = isapprox(a.x, b.x; kwargs...)
 
 gentype(μ::Dirac{X}) where {X} = X
diff --git a/src/transport.jl b/src/transport.jl
index ba8387be..1eed0e62 100644
--- a/src/transport.jl
+++ b/src/transport.jl
@@ -90,6 +90,7 @@ end
 function Base.:(==)(a::TransportFunction, b::TransportFunction)
     return a.ν == b.ν && a.μ == b.μ
 end
+Base.hash(f::TransportFunction, h::UInt) = hash(f.ν, hash(f.μ, hash(:TransportFunction, h)))
 
 Base.@propagate_inbounds function (f::TransportFunction)(x)
     return transport_def(f.ν, f.μ, checked_arg(f.μ, x))
diff --git a/test/test_basics.jl b/test/test_basics.jl
index 8dbdbe4d..71540e22 100644
--- a/test/test_basics.jl
+++ b/test/test_basics.jl
@@ -224,6 +224,14 @@ end
     @test MeasureBase.mspace_elsize(Dirac([1, 2])) == (2,)
 end
 
+@testset "hash follows equality" begin
+    a = productmeasure([Dirac([1.0, 2.0]), StdNormal()])
+    b = productmeasure([Dirac([1.0, 2.0]), StdNormal()])
+    @test a == b && hash(a) == hash(b)
+    f, g = transport_to(StdUniform()^1, a), transport_to(StdUniform()^1, b)
+    @test f == g && hash(f) == hash(g)
+end
+
 @testset "logdensity_rel" begin
     @test logdensity_rel(Dirac(0.0) + Lebesgue(), Dirac(1.0), 0.0) == Inf
     @test logdensity_rel(Dirac(0.0) + Lebesgue(), Dirac(1.0), 1.0) == -Inf

From ea8fda08ef4306c9e6200a44abd527b5a6cc207a Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sat, 19 Sep 2026 14:33:18 +0200
Subject: [PATCH 110/122] Learn pushforward sizes from every fixed-layout
 origin

The output size of a pushforward is learned from a test value whenever
the origin's variates have a fixed layout, not only when the origin has
a flat size. Pushforwards of tuple products (such as unshaped measures)
thereby declare their variate rank, so their batched kernels and the
transport broadcast hook no longer fall back to pointwise evaluation.

Created by generative AI.
---
 src/combinators/transformedmeasure.jl  | 10 ++++++----
 test/combinators/transformedmeasure.jl |  9 +++++++++
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 72719767..b0faffa5 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -86,11 +86,13 @@ end
 
 # The size of the variates of a pushforward follows from a test value of
 # the origin, where the origin has variates of known size:
-@inline function _pushfwd_varsize(f, μ)
-    _pushfwd_varsize(f, μ, mspace_flatsize(μ))
+# The output size is learned from a test value whenever the origin's
+# variates have a fixed layout, which includes tuple products:
+@inline function _pushfwd_varsize(f, μ::MU) where {MU}
+    _pushfwd_varsize(f, μ, fixed_stream_size(MU))
 end
-@inline _pushfwd_varsize(f, μ, ::SizeLike) = _value_flatsize(f(testvalue(μ)))
-@inline _pushfwd_varsize(f, μ, ::NoMSpaceElementSize) = NoMSpaceElementSize{typeof(μ)}()
+@inline _pushfwd_varsize(f, μ, ::True) = _value_flatsize(f(testvalue(μ)))
+@inline _pushfwd_varsize(f, μ, ::False) = NoMSpaceElementSize{typeof(μ)}()
 
 @inline mspace_elsize(ν::PushforwardMeasure) = _value_or_unknown(ν.varsize, ν)
 @inline mspace_flatsize(ν::PushforwardMeasure) = _value_or_unknown(ν.varsize, ν)
diff --git a/test/combinators/transformedmeasure.jl b/test/combinators/transformedmeasure.jl
index 6261c762..2d47900f 100644
--- a/test/combinators/transformedmeasure.jl
+++ b/test/combinators/transformedmeasure.jl
@@ -4,6 +4,7 @@ using MeasureBase
 using MeasureBase: pushfwd, StdUniform, StdExponential, StdLogistic
 using MeasureBase: pushfwd, PushforwardMeasure
 using MeasureBase: transport_to, unsafe_logdensityof
+using MeasureBase: productmeasure, mbind
 import Statistics: var
 using DensityInterface: logdensityof
 using LogExpFunctions
@@ -179,4 +180,12 @@ end
     @test PushfwdRootMeasure() isa PushFwdStyle
     @test MeasureBase.WithVolCorr === AdaptRootMeasure
     @test MeasureBase.NoVolCorr === PushfwdRootMeasure
+
+    @testset "output size of pushforwards of tuple products" begin
+        Pt = productmeasure((StdNormal(), StdUniform()^2))
+        ν = pushfwd(x -> vcat(x[1], x[2]), Pt)
+        @test MeasureBase.mspace_flatsize(ν) == (3,)
+        @test MeasureBase.mspace_ndims(typeof(ν)) == 1
+        @test MeasureBase.mspace_flatsize(pushfwd(x -> x, mbind(x -> StdNormal()^(x > 0 ? 1 : 2), StdNormal()))) isa MeasureBase.NoMSpaceElementSize
+    end
 end

From 44afb4d809e6f38bae3534d4cefb102469677424 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sat, 19 Sep 2026 14:52:14 +0200
Subject: [PATCH 111/122] Fold fixed stream sizes of tuple products to
 constants

The check whether all marginals of a tuple product have fixed stream
sizes is folded pairwise over the marginal types, so that it infers as
a constant and the types of pushforwards of such products stay
inferrable.

Created by generative AI.
---
 src/combinators/product.jl | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index e370e62a..f90d065d 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -688,9 +688,14 @@ function batched_logdensityof_with_rest(μ::ProductMeasure{<:NamedTuple{names}},
     batched_logdensityof_with_rest(productmeasure(values(marginals(μ))), x, sz)
 end
 
-@inline function fixed_stream_size(::Type{<:ProductMeasure{M}}) where {M<:Tuple}
-    static(all(T -> fixed_stream_size(T) === static(true), M.parameters))
-end
+# Folded pairwise over the marginal types, so that the result is a constant:
+@inline fixed_stream_size(::Type{<:ProductMeasure{M}}) where {M<:Tuple} = _all_fixed_stream_sizes(M)
+@inline _all_fixed_stream_sizes(::Type{Tuple{}}) = static(true)
+@inline function _all_fixed_stream_sizes(::Type{M}) where {M<:Tuple}
+    _both(fixed_stream_size(Base.tuple_type_head(M)), _all_fixed_stream_sizes(Base.tuple_type_tail(M)))
+end
+@inline _both(::True, ::True) = static(true)
+@inline _both(::StaticBool, ::StaticBool) = static(false)
 @inline function fixed_stream_size(::Type{<:ProductMeasure{NamedTuple{names,M}}}) where {names,M<:Tuple}
     fixed_stream_size(ProductMeasure{M})
 end

From 236d23fe4e1aae1411318e5b7f66b458262c2292 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Sat, 19 Sep 2026 15:01:46 +0200
Subject: [PATCH 112/122] Keep variate layouts out of autodiff

The fixed-stream-size check and the output size of pushforwards are
type-level information, declared non-differentiable so that Zygote can
differentiate through the construction of pushforwards of tuple
products.

Created by generative AI.
---
 ext/MeasureBaseChainRulesCoreExt.jl    |  7 ++++++-
 test/combinators/transformedmeasure.jl | 16 ++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/ext/MeasureBaseChainRulesCoreExt.jl b/ext/MeasureBaseChainRulesCoreExt.jl
index cd042496..f78adefd 100644
--- a/ext/MeasureBaseChainRulesCoreExt.jl
+++ b/ext/MeasureBaseChainRulesCoreExt.jl
@@ -3,7 +3,7 @@
 module MeasureBaseChainRulesCoreExt
 
 using MeasureBase
-using ChainRulesCore: NoTangent, ZeroTangent
+using ChainRulesCore: NoTangent, ZeroTangent, @non_differentiable
 import ChainRulesCore
 
 # = utils ====================================================================
@@ -73,6 +73,11 @@ end
 
 # = insupport & friends ======================================================
 
+# Variate layouts are type-level information:
+using MeasureBase: fixed_stream_size, _pushfwd_varsize
+@non_differentiable fixed_stream_size(::Type)
+@non_differentiable _pushfwd_varsize(f, μ)
+
 using MeasureBase: check_dof, require_insupport, checked_arg, _checksupport
 
 @inline function ChainRulesCore.rrule(::typeof(_checksupport), cond, result)
diff --git a/test/combinators/transformedmeasure.jl b/test/combinators/transformedmeasure.jl
index 2d47900f..a7f80ac5 100644
--- a/test/combinators/transformedmeasure.jl
+++ b/test/combinators/transformedmeasure.jl
@@ -5,6 +5,9 @@ using MeasureBase: pushfwd, StdUniform, StdExponential, StdLogistic
 using MeasureBase: pushfwd, PushforwardMeasure
 using MeasureBase: transport_to, unsafe_logdensityof
 using MeasureBase: productmeasure, mbind
+import Zygote
+using InverseFunctions: inverse
+using ChangesOfVariables: with_logabsdet_jacobian
 import Statistics: var
 using DensityInterface: logdensityof
 using LogExpFunctions
@@ -188,4 +191,17 @@ end
         @test MeasureBase.mspace_ndims(typeof(ν)) == 1
         @test MeasureBase.mspace_flatsize(pushfwd(x -> x, mbind(x -> StdNormal()^(x > 0 ? 1 : 2), StdNormal()))) isa MeasureBase.NoMSpaceElementSize
     end
+
+    @testset "construction inside differentiated functions" begin
+        flat(x::Tuple) = vcat(x[1], x[2])
+        unflat(v::AbstractVector) = (v[1], v[2:end])
+        InverseFunctions.inverse(::typeof(flat)) = unflat
+        InverseFunctions.inverse(::typeof(unflat)) = flat
+        ChangesOfVariables.with_logabsdet_jacobian(::typeof(flat), x) = (flat(x), zero(eltype(x[2])))
+        ChangesOfVariables.with_logabsdet_jacobian(::typeof(unflat), v) = (unflat(v), zero(eltype(v)))
+        Pt = productmeasure((StdNormal(), StdUniform()^2))
+        g(v) = logdensityof(pushfwd(flat, Pt), v)
+        v = [0.3, 0.2, 0.7]
+        @test Zygote.gradient(g, v)[1] ≈ [-0.3, 0.0, 0.0]
+    end
 end

From 6cbdbd4d1d8d69bf9b1a74e70e763ee65d31d154 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 21 Sep 2026 14:40:48 +0200
Subject: [PATCH 113/122] Keep finite inputs finite in transports

On the floating-point grid the endpoints of the unit interval stand for
their nearest interior points: uniform inputs are clamped into the open
interval (from the smallest normal float, since devices may flush
subnormals, to the grid point below one) before quantiles, and tail
probabilities of the direct log-space conversions between standard
measures never underflow to zero. Inputs outside the unit interval stay
NaN. This is the null-set convention BAT has used in practice, applied
once in MeasureBase for the standard measures, Half, the wrapped
families, Dirichlet and the generic logistic path; the bounds are
non-differentiable helpers so that autodiff passes through them.

Created by generative AI.
---
 ext/MeasureBaseChainRulesCoreExt.jl           |  4 ++-
 .../MeasureBaseDistributionsExt.jl            |  1 +
 ext/MeasureBaseDistributionsExt/families.jl   |  8 +++---
 .../multivariate.jl                           |  2 +-
 ext/MeasureBaseDistributionsExt/univariate.jl |  2 +-
 redesign.md                                   | 25 +++++++++++--------
 src/combinators/half.jl                       |  2 +-
 src/density-core.jl                           | 10 ++++++++
 src/standard/stdconvert.jl                    |  6 ++---
 src/standard/stdexponential.jl                |  2 +-
 src/standard/stdlogistic.jl                   |  4 +--
 src/standard/stdnormal.jl                     |  6 ++---
 src/transport.jl                              | 13 +++++++---
 test/support_conventions.jl                   | 12 +++++++--
 14 files changed, 63 insertions(+), 34 deletions(-)

diff --git a/ext/MeasureBaseChainRulesCoreExt.jl b/ext/MeasureBaseChainRulesCoreExt.jl
index f78adefd..ddac715d 100644
--- a/ext/MeasureBaseChainRulesCoreExt.jl
+++ b/ext/MeasureBaseChainRulesCoreExt.jl
@@ -74,9 +74,11 @@ end
 # = insupport & friends ======================================================
 
 # Variate layouts are type-level information:
-using MeasureBase: fixed_stream_size, _pushfwd_varsize
+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
 
diff --git a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
index 7795e25a..c0ca18f0 100644
--- a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
+++ b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
@@ -15,6 +15,7 @@ import MeasureBase
 using MeasureBase: AbstractMeasure, AsMeasure, asmeasure
 using MeasureBase: Lebesgue, Counting, ℝ
 using MeasureBase: StdMeasure, StdUniform, StdExponential, StdLogistic, StdNormal
+using MeasureBase: _unit_interior
 using MeasureBase: PowerMeasure, WeightedMeasure, SuperpositionMeasure, PushforwardMeasure
 using MeasureBase: basemeasure, rootmeasure, testvalue, productmeasure, pushfwd, superpose
 using MeasureBase: getdof, checked_arg, massof
diff --git a/ext/MeasureBaseDistributionsExt/families.jl b/ext/MeasureBaseDistributionsExt/families.jl
index ffd24a9c..54a29633 100644
--- a/ext/MeasureBaseDistributionsExt/families.jl
+++ b/ext/MeasureBaseDistributionsExt/families.jl
@@ -105,7 +105,7 @@ end
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Cauchy}) = StdUniform
 @inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Cauchy, x) = 1 // 2 + atan((x - d.μ) / d.σ) / π
 @inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Cauchy, p) =
-    _nan_outside(StdUniform, p, muladd(d.σ, tan(π * (p - 1 // 2)), d.μ))
+    _nan_outside(StdUniform, p, muladd(d.σ, tan(π * (_unit_interior(p) - 1 // 2)), d.μ))
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Laplace}) = StdUniform
 @inline function MeasureBase.transport_to_std(::Type{StdUniform}, d::Laplace, x)
@@ -114,7 +114,7 @@ end
 end
 @inline function MeasureBase.transport_from_std(::Type{StdUniform}, d::Laplace, p)
     u = p - 1 // 2
-    _nan_outside(StdUniform, p, muladd(-d.θ * sign(u), log1p(-min(2 * abs(u), one(u))), d.μ))
+    _nan_outside(StdUniform, p, muladd(-d.θ * sign(u), log1p(-abs(2 * _unit_interior(p) - 1)), d.μ))
 end
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:LogNormal}) = StdNormal
@@ -127,10 +127,10 @@ end
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Gamma}) = StdUniform
 @inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Gamma, x) = _nan_outside(d, x, _gamma_cdf(d.α, abs(x / d.θ)))
-@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Gamma, p) = _nan_outside(StdUniform, p, d.θ * _gamma_quantile(d.α, _unit_clamp(p)))
+@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Gamma, p) = _nan_outside(StdUniform, p, d.θ * _gamma_quantile(d.α, _unit_interior(p)))
 
 @inline MeasureBase.preferred_stdmeasure(::Type{<:Beta}) = StdUniform
 @inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Beta, x) = _nan_outside(d, x, _beta_cdf(d.α, d.β, _unit_clamp(x)))
-@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Beta, p) = _nan_outside(StdUniform, p, _beta_quantile(d.α, d.β, _unit_clamp(p)))
+@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Beta, p) = _nan_outside(StdUniform, p, _beta_quantile(d.α, d.β, _unit_interior(p)))
 
 @inline _unit_clamp(x) = clamp(x, zero(x), one(x))
diff --git a/ext/MeasureBaseDistributionsExt/multivariate.jl b/ext/MeasureBaseDistributionsExt/multivariate.jl
index b4262882..1ddfe237 100644
--- a/ext/MeasureBaseDistributionsExt/multivariate.jl
+++ b/ext/MeasureBaseDistributionsExt/multivariate.jl
@@ -108,7 +108,7 @@ function MeasureBase.batched_transport_from_std(::Type{StdUniform}, d::Dirichlet
     K = length(d)
     αs, βs = _stick_breaking_params(d)
     Zc = _as_columns(Z)
-    beta_v = _beta_quantile.(αs, βs, _unit_clamp.(Zc))
+    beta_v = _beta_quantile.(αs, βs, _unit_interior.(Zc))
     cp = cumprod(beta_v; dims = 1)
     # Each component takes what its Beta variate leaves of the remaining mass:
     X = vcat(1 .- _rows(cp, 1:1), _rows(cp, 1:(K - 2)) .- _rows(cp, 2:(K - 1)), _rows(cp, (K - 1):(K - 1)))
diff --git a/ext/MeasureBaseDistributionsExt/univariate.jl b/ext/MeasureBaseDistributionsExt/univariate.jl
index 1f318379..27d5f523 100644
--- a/ext/MeasureBaseDistributionsExt/univariate.jl
+++ b/ext/MeasureBaseDistributionsExt/univariate.jl
@@ -77,7 +77,7 @@ end
 @inline function MeasureBase.transport_from_std(::Type{StdLogistic}, d::Distribution{Univariate,Continuous}, l)
     R = _result_numtype(d, l)
     # From the side that keeps the tail:
-    x = l < zero(l) ? _trafo_quantile(d, logistic(l)) : _trafo_cquantile(d, logistic(-l))
+    x = l < zero(l) ? _trafo_quantile(d, _unit_interior(logistic(l))) : _trafo_cquantile(d, _unit_interior(logistic(-l)))
     convert(R, x)
 end
 
diff --git a/redesign.md b/redesign.md
index be52602e..90302dd5 100644
--- a/redesign.md
+++ b/redesign.md
@@ -189,20 +189,25 @@ locally on the GB10, green at HEAD except one expected-broken CUDA case
   `checked_arg` at the entry points; variates of the right shape never
   throw: densities are `-Inf` outside the support (including non-integers
   for counting-based measures and infinite values), transports are `NaN`
-  outside the support of the source (boundaries may map to `±Inf`), `NaN`
-  inputs give `NaN` or `-Inf`, relative densities keep `+Inf`/`-Inf`/`NaN`.
+  outside the support of the source, `NaN` inputs give `NaN` or `-Inf`,
+  relative densities keep `+Inf`/`-Inf`/`NaN`. Finite inputs give finite
+  transports: uniform inputs are clamped into the open unit interval
+  before quantiles (the endpoints stand for their nearest interior grid
+  points, a null-set convention like BAT's `[eps, 1 - eps]` clamping in
+  practice), tail probabilities of log-space conversions floor at the
+  smallest positive float.
   Kernels must not throw outside the support, since the masks evaluate
   both branches (`abs`, `clamp`, `min` guards instead of `NaNMath`, which
   isn't device-compatible). Downstream checks such as BAT's
   `checked_logdensityof` stay downstream.
-- BAT's boundary tweaks (to discuss): adopted are infinite variates
-  outside the support of continuous wrapped distributions, `-Inf + Inf`
-  giving `-Inf` in pushforward densities and a zero Jacobian term where
-  both densities vanish. Not adopted: clamping uniform-direction inputs to
-  `[eps, 1 - eps]` (here `NaN` outside, `±Inf` or support edges at 0 and
-  1), replacing finite densities with infinite Jacobian terms by `-1e38`,
-  and re-evaluating densities an `eps` inside the support where
-  Distributions returns `NaN` (the family kernels are exact there).
+- BAT's boundary tweaks: adopted are infinite variates outside the
+  support of continuous wrapped distributions, the clamping of uniform
+  inputs into the open unit interval (for inputs inside `[0, 1]`, outside
+  stays `NaN`), `-Inf + Inf` giving `-Inf` in pushforward densities and a zero
+  Jacobian term where both densities vanish. Not adopted: replacing finite
+  densities with infinite Jacobian terms by `-1e38`, and re-evaluating
+  densities an `eps` inside the support where Distributions returns `NaN`
+  (the family kernels are exact there).
   Quantile results within `4 eps` of the support edges snap to the edges
   only on the generic logistic path of wrapped distributions.
 - Decisions pending: `Half` tails via log-ccdf, device random variate
diff --git a/src/combinators/half.jl b/src/combinators/half.jl
index f793a947..20ee71c0 100644
--- a/src/combinators/half.jl
+++ b/src/combinators/half.jl
@@ -46,7 +46,7 @@ function smf(μ::Half, x)
 end
 
 function invsmf(μ::Half, p)
-    _nan_outside(StdUniform(), p, invsmf(μ.parent, (min(p, one(p)) + 1) / 2))
+    _nan_outside(StdUniform(), p, invsmf(μ.parent, _unit_interior((p + 1) / 2)))
 end
 
 @inline transport_to_std(::Type{StdUniform}, μ::Half, x) = _nan_outside(μ, x, smf(μ, x))
diff --git a/src/density-core.jl b/src/density-core.jl
index e39f9809..3859f28a 100644
--- a/src/density-core.jl
+++ b/src/density-core.jl
@@ -87,6 +87,16 @@ end
 # support:
 @inline _nan_outside(μ, x, y) = ifelse(_insupport_mask(insupport(μ, x)), y, oftype(y, NaN))
 
+# On the floating-point grid the endpoints of the unit interval stand for
+# their nearest interior points (the smallest normal float above zero,
+# since devices may flush subnormals, and the grid point below one), so
+# that quantiles stay finite, and tail probabilities in log-space
+# conversions never underflow to zero:
+@inline _unit_interior(p) = clamp(p, _unit_bounds(p)...)
+@inline _unit_bounds(p) = (_prob_floor(p), prevfloat(one(p)))
+@inline _positive_prob(p) = max(p, _prob_floor(p))
+@inline _prob_floor(p) = floatmin(typeof(one(p)))
+
 """
     MeasureBase.logdensityof_with_rest(μ::AbstractMeasure, x)
 
diff --git a/src/standard/stdconvert.jl b/src/standard/stdconvert.jl
index 3868684b..520f3097 100644
--- a/src/standard/stdconvert.jl
+++ b/src/standard/stdconvert.jl
@@ -10,17 +10,17 @@
 @inline _normccdf(z) = erfc(z * invsqrt2) / 2
 
 @inline function transport_def(::StdExponential, ::StdNormal, z)
-    ifelse(z < zero(z), -log1p(-Φ(z)), -log(_normccdf(z)))
+    ifelse(z < zero(z), -log1p(-Φ(z)), -log(_positive_prob(_normccdf(z))))
 end
 
 @inline function transport_def(::StdNormal, μ::StdExponential, x)
-    _nan_outside(μ, x, ifelse(x < oftype(x, logtwo), Φinv(-expm1(-x)), -Φinv(exp(-x))))
+    _nan_outside(μ, x, ifelse(x < oftype(x, logtwo), Φinv(_positive_prob(-expm1(-x))), -Φinv(min(_positive_prob(exp(-x)), one(x)))))
 end
 
 @inline transport_def(::StdLogistic, ::StdNormal, z) = _normlogcdf(z) - _normlogccdf(z)
 
 @inline function transport_def(::StdNormal, ::StdLogistic, l)
-    ifelse(l < zero(l), Φinv(logistic(l)), -Φinv(logistic(-l)))
+    ifelse(l < zero(l), Φinv(_positive_prob(logistic(l))), -Φinv(_positive_prob(logistic(-l))))
 end
 
 @inline transport_def(::StdLogistic, μ::StdExponential, x) = _nan_outside(μ, x, log(abs(expm1(-x))) + x)
diff --git a/src/standard/stdexponential.jl b/src/standard/stdexponential.jl
index 1bb89c24..a9823084 100644
--- a/src/standard/stdexponential.jl
+++ b/src/standard/stdexponential.jl
@@ -19,7 +19,7 @@ end
 @inline basemeasure(::StdExponential) = LebesgueBase()
 
 @inline transport_def(::StdUniform, μ::StdExponential, x) = _nan_outside(μ, x, -expm1(-x))
-@inline transport_def(::StdExponential, μ::StdUniform, x) = _nan_outside(μ, x, -log1p(-min(x, one(x))))
+@inline transport_def(::StdExponential, μ::StdUniform, x) = _nan_outside(μ, x, -log1p(-_unit_interior(x)))
 
 @inline rand_impl(ctx::GenContext, ::StdExponential) = randexp(get_rng(ctx), get_precision(ctx))
 @inline batched_rand_impl(ctx::GenContext, ::StdExponential, sz::Dims) = _randexp_bulk(ctx, sz)
diff --git a/src/standard/stdlogistic.jl b/src/standard/stdlogistic.jl
index d1cfcd44..98b76219 100644
--- a/src/standard/stdlogistic.jl
+++ b/src/standard/stdlogistic.jl
@@ -16,9 +16,7 @@ export StdLogistic
 @inline basemeasure(::StdLogistic) = LebesgueBase()
 
 @inline transport_def(::StdUniform, μ::StdLogistic, x) = logistic(x)
-@inline transport_def(::StdLogistic, μ::StdUniform, p) = _nan_outside(μ, p, _logit_nan(p))
-# Total on the CPU, the mask supplies the NaN outside of [0, 1]:
-@inline _logit_nan(p) = log(abs(p) / abs(one(p) - p))
+@inline transport_def(::StdLogistic, μ::StdUniform, p) = _nan_outside(μ, p, logit(_unit_interior(p)))
 
 @inline rand_impl(ctx::GenContext, ::StdLogistic) = logit(rand(get_rng(ctx), get_precision(ctx)))
 @inline batched_rand_impl(ctx::GenContext, ::StdLogistic, sz::Dims) = logit.(_rand_bulk(ctx, sz))
diff --git a/src/standard/stdnormal.jl b/src/standard/stdnormal.jl
index 21312653..58b0408e 100644
--- a/src/standard/stdnormal.jl
+++ b/src/standard/stdnormal.jl
@@ -24,16 +24,16 @@ export StdNormal
 @inline batched_rand_impl(ctx::GenContext, ::StdNormal, sz::Dims) = _randn_bulk(ctx, sz)
 
 Φ(z) = erfc(-z * invsqrt2) / 2
-Φinv(p) = -erfcinv(2 * clamp(p, zero(p), one(p))) * sqrt2
+Φinv(p) = -erfcinv(2 * p) * sqrt2
 
 InverseFunctions.inverse(::typeof(Φ)) = Φinv
 InverseFunctions.inverse(::typeof(Φinv)) = Φ
 
 smf(::StdNormal, x) = Φ(x)
-invsmf(::StdNormal, p) = Φinv(p)
+invsmf(::StdNormal, p) = _nan_outside(StdUniform(), p, Φinv(_unit_interior(p)))
 
 smf(::StdNormal) = Φ
 invsmf(::StdNormal) = Φinv
 
-transport_def(::StdNormal, μ::StdUniform, p) = _nan_outside(μ, p, Φinv(p))
+transport_def(::StdNormal, μ::StdUniform, p) = _nan_outside(μ, p, Φinv(_unit_interior(p)))
 transport_def(::StdUniform, ::StdNormal, x) = Φ(x)
diff --git a/src/transport.jl b/src/transport.jl
index 1eed0e62..4c7403c0 100644
--- a/src/transport.jl
+++ b/src/transport.jl
@@ -56,10 +56,15 @@ Transport `x` from the measure `μ` to the measure `ν`, equivalent to
 # Extended help
 
 Variates of the right shape never throw: outside the support of `μ` the
-result is `NaN` (elementwise for powers and products), values at the
-boundary of the support may map to infinite values. Variates of the wrong
-shape throw an `ArgumentError`. Transport implementations must not throw
-outside the support, since the `NaN` masks evaluate both branches.
+result is `NaN` (elementwise for powers and products). Variates of the
+wrong shape throw an `ArgumentError`. Transport implementations must not
+throw outside the support, since the `NaN` masks evaluate both branches.
+
+Finite inputs give finite results: on the floating-point grid the
+endpoints of the unit interval stand for their nearest interior grid
+points (uniform inputs are clamped into the open interval before
+quantiles), and tail probabilities in log-space conversions never
+underflow to zero.
 """
 transport_to(ν, μ, x) = transport_to(ν, μ)(x)
 
diff --git a/test/support_conventions.jl b/test/support_conventions.jl
index 9cea29c4..72b36837 100644
--- a/test/support_conventions.jl
+++ b/test/support_conventions.jl
@@ -45,8 +45,16 @@ using JLArrays: JLArray
             @test isnan(transport_to(ν, StdUniform())(1.5))
             @test isnan(transport_to(ν, StdUniform())(-0.5))
         end
-        @test transport_to(StdExponential(), StdUniform())(1.0) == Inf
-        @test transport_to(StdNormal(), StdUniform())(0.0) == -Inf
+        # Endpoints of the unit interval stand for their nearest interior
+        # points, tails never underflow to infinite variates:
+        for ν in (StdExponential(), StdLogistic(), StdNormal(), Half(StdNormal()))
+            f = transport_to(ν, StdUniform())
+            @test isfinite(f(0.0)) && isfinite(f(1.0)) && f(0.0) <= f(0.5) <= f(1.0)
+            @test f(1.0) == f(prevfloat(1.0)) && f(0.0) == f(floatmin(Float64))
+        end
+        for (ν, μ) in ((StdExponential(), StdNormal()), (StdNormal(), StdExponential()), (StdNormal(), StdLogistic()))
+            @test all(isfinite, transport_to(ν, μ).([-1e6, -40.0, 40.0, 1e6][MeasureBase.insupport.(Ref(μ), [-1e6, -40.0, 40.0, 1e6])]))
+        end
         @test !isnan(transport_to(StdUniform(), StdNormal())(-37.0))
         @test !isnan(transport_to(StdUniform(), StdLogistic())(-800.0))
     end

From 1b1a32fcc10da76e76a060d5d48fca2b08c9caa8 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 21 Sep 2026 16:25:15 +0200
Subject: [PATCH 114/122] Build the static helpers on StaticThings

The private static helpers duplicated tools that StaticThings now
provides: `_size_dims`, `_reshape_batch`, the leading-dimension family
(`_sum_leading_dims`, `_drop_leading_dims`, `_merge_leading_dims`,
`_all_leading_dims`), `_get_or_view`/`_split_after` for vectors, the
type-level `_static_axes_size`, and the pairwise folds over tuple types
of products and superpositions. They are gone in favour of `size_dims`,
`maybestatic_reshape`, `sum_leading_dims`, `drop_leading_dims`,
`merge_leading_dims`, `all_leading_dims`, `maybestatic_view`, `split_at`,
`axes2size(::Type)`, `static_all` and `static_reduce`. Hand-rolled
`map(dynamic, ...)` became `asnonstatic` and `_both` Static's `&`.

`mspace_ndims` now follows from `maybestatic_length` and returns an
`IntegerLike`. `Static.StaticInt` is not an `Integer`, so the dispatch
that consumes it takes `IntegerLike` and `_unit_dof` compares static
ranks instead of testing for a literal zero.

Requires StaticThings 0.3, whose `maybestatic_reshape` no longer turns
arbitrary arrays into static arrays.

Created by generative AI.
---
 Project.toml                          |  2 +-
 src/MeasureBase.jl                    | 17 +++---
 src/collection_utils.jl               | 32 ++----------
 src/combinators/bind.jl               |  6 +--
 src/combinators/combined.jl           |  8 +--
 src/combinators/power.jl              | 40 +++++++-------
 src/combinators/product.jl            | 18 +++----
 src/combinators/reshape.jl            |  2 +-
 src/combinators/superpose.jl          | 28 ++++------
 src/combinators/transformedmeasure.jl |  2 +-
 src/density-batched.jl                | 75 ++++++---------------------
 src/mspace.jl                         | 12 ++---
 src/primitives/dirac.jl               |  7 +--
 src/standard/stdconvert.jl            |  2 +-
 src/transport-batched.jl              | 23 +++-----
 src/transport.jl                      |  2 +-
 16 files changed, 87 insertions(+), 189 deletions(-)

diff --git a/Project.toml b/Project.toml
index a8e8b924..7ea88557 100644
--- a/Project.toml
+++ b/Project.toml
@@ -100,7 +100,7 @@ Reexport = "1"
 SpecialFunctions = "2.1.4"
 Static = "0.8, 1"
 StaticArrays = "1.5"
-StaticThings = "0.2"
+StaticThings = "0.3"
 Statistics = "1"
 StructArrays = "0.6.18, 0.7"
 StatsBase = "0.33, 0.34"
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 8d5dafb9..976a6731 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -45,17 +45,14 @@ using FunctionChains
 using PropertyFunctions: PropSelFunction
 
 using StaticThings:
-    AxesLike, StaticAxesLike, SizeLike, StaticSizeLike,
-    OneToLike, StaticOneTo, StaticOneToLike, RealLike,
-    IntegerLike, StaticUnitRange, StaticUnitRangeLike,
-    NoTypeSize,
-    asaxes, asnonstatic, 
-    canonical_axes, canonical_indices, canonical_size,
-    maybestatic_axes, maybestatic_eachindex,
-    maybestatic_length, maybestatic_size, maybestatic_first, maybestatic_last,
+    SizeLike, OneToLike, StaticOneToLike, IntegerLike,
+    asaxes, asnonstatic, canonical_size, size_dims,
+    maybestatic_eachindex, maybestatic_length, maybestatic_size,
+    maybestatic_first, maybestatic_last, maybestatic_view,
     maybestatic_oneto, maybestatic_fill, maybestatic_reshape,
-    size_from_type, axes2size, size2axes, size2length,
-    staticarray_type
+    axes2size, size2length, split_at,
+    static_all, static_any, static_reduce,
+    sum_leading_dims, drop_leading_dims, merge_leading_dims, all_leading_dims
 
 import HeterogeneousComputing
 using HeterogeneousComputing: real_numtype
diff --git a/src/collection_utils.jl b/src/collection_utils.jl
index 009bfe00..dbf66acc 100644
--- a/src/collection_utils.jl
+++ b/src/collection_utils.jl
@@ -26,30 +26,6 @@ _exp_cumsum_log(xs::AbstractVector) = exp.(cumsum(log.(xs)))
 Base.@propagate_inbounds _as_tuple(v::AbstractVector, ::Val{N}) where {N} = Tuple(SVector{N}(v))
 
 
-Base.@propagate_inbounds function _get_or_view(A::AbstractVector, from::IntegerLike, until::IntegerLike)
-    view(A, dynamic(from):dynamic(until))
-end
-
-Base.@propagate_inbounds function _get_or_view(
-    A::StaticVector,
-    from::StaticInteger{F},
-    until::StaticInteger{U},
-) where {F,U}
-    SVector{U - F + 1,eltype(A)}(_get_or_view(Tuple(A), from, until))
-end
-
-Base.@propagate_inbounds function _get_or_view(tpl::Tuple, from::IntegerLike, until::IntegerLike)
-    ntuple(i -> tpl[from + i - 1], Val(until - from + 1))
-end
-
-
-@inline function _split_after(x::AbstractVector, n::IntegerLike)
-    idxs = maybestatic_eachindex(x)
-    i_first = maybestatic_first(idxs)
-    i_last = maybestatic_last(idxs)
-    _get_or_view(x, i_first, i_first + n - one(n)), _get_or_view(x, i_first + n, i_last)
-end
-
 @inline _split_after(x::Tuple, n) = _split_after(x::Tuple, Val{n}())
 @inline _split_after(x::Tuple, ::Val{N}) where {N} = x[begin:(begin+N-1)], x[(begin+N):end]
 
@@ -96,21 +72,21 @@ _cat_measures(a::AbstractVector, b::AbstractVector) = vcat(a, b)
 # Take the beginning of a flat vector stream as a variate of size `sz`,
 # scalar variates have size `()` and multi-rank variates are reshaped:
 Base.@propagate_inbounds _consume_from_stream(x::AbstractVector, sz::Tuple{IntegerLike}) =
-    _split_after(x, sz[1])
+    split_at(x, sz[1])
 
 Base.@propagate_inbounds function _consume_from_stream(x::AbstractVector, ::Tuple{})
     idxs = maybestatic_eachindex(x)
     i_first = maybestatic_first(idxs)
-    x[i_first], _get_or_view(x, i_first + one(i_first), maybestatic_last(idxs))
+    x[i_first], maybestatic_view(x, i_first + one(i_first), maybestatic_last(idxs))
 end
 
 function _consume_from_stream(x::AbstractVector, sz::Tuple{Vararg{IntegerLike}})
-    a_flat, x_rest = _split_after(x, size2length(sz))
+    a_flat, x_rest = split_at(x, size2length(sz))
     return maybestatic_reshape(a_flat, sz), x_rest
 end
 
 Base.@propagate_inbounds _consume_from_stream(x::AbstractVector, sz::StaticArrays.Size) =
-    _consume_from_stream(x, _size_dims(sz))
+    _consume_from_stream(x, size_dims(sz))
 
 function _consume_from_stream(x::AbstractVector, @nospecialize(sz))
     throw(ArgumentError("Can't consume a variate of size $sz from a flat vector stream"))
diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl
index 87c43bac..39047fc7 100644
--- a/src/combinators/bind.jl
+++ b/src/combinators/bind.jl
@@ -259,7 +259,7 @@ end
 function _bind_tpm_sc_cat(f_c::typeof(vcat), μ::_BindBy{typeof(vcat)}, xy::AbstractVector)
     tpm_μ, a, b, y, xy = _bind_tpm_sc_cat_lμabyxy(f_c, μ, xy)
     # Don't use `x = f_c(a, b)` here, would allocate, splitting xy can use views:
-    x, y = _split_after(xy, length(a) + length(b))
+    x, y = split_at(xy, maybestatic_length(a) + maybestatic_length(b))
     return tpm_μ, x, y
 end
 
@@ -330,7 +330,7 @@ end
 function logdensityof_with_rest(μ::_BindBy{typeof(vcat)}, x::AbstractVector)
     ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x)
     ℓ_b, b, x_rest = logdensityof_with_rest(_get_β_a(μ, a), x2)
-    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest))
     return ℓ_a + ℓ_b, x_μ, x_rest
 end
 
@@ -387,7 +387,7 @@ end
 function transport_to_std_with_rest(::Type{S}, μ::_BindBy{typeof(vcat)}, x::AbstractVector) where {S<:StdMeasure}
     z_a, a, x2 = transport_to_std_with_rest(S, μ.α, x)
     z_b, _, x_rest = transport_to_std_with_rest(S, _get_β_a(μ, a), x2)
-    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest))
     return vcat(z_a, z_b), x_μ, x_rest
 end
 
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index fdfe07bd..b5c54d12 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -31,7 +31,7 @@ function _generic_split_combined(f_c::FC, α::AbstractMeasure, ab) where {FC}
 end
 
 _split_variate_byvalue(::typeof(vcat), test_a::AbstractVector, ab::AbstractVector) =
-    _split_after(ab, length(test_a))
+    split_at(ab, maybestatic_length(test_a))
 
 _split_variate_byvalue(::typeof(vcat), ::Number, ab::AbstractVector) =
     _consume_from_stream(ab, ())
@@ -186,7 +186,7 @@ end
 
 @inline mspace_ndims(::Type{<:CombinedMeasure{typeof(vcat)}}) = 1
 @inline function fixed_stream_size(::Type{<:CombinedMeasure{<:Any,MA,MB}}) where {MA,MB}
-    static(fixed_stream_size(MA) === static(true) && fixed_stream_size(MB) === static(true))
+    fixed_stream_size(MA) & fixed_stream_size(MB)
 end
 
 # Batches of vcat-combined variates are batches of streams: with fixed
@@ -234,7 +234,7 @@ end
 function logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector)
     ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x)
     ℓ_b, b, x_rest = logdensityof_with_rest(μ.β, x2)
-    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest))
     return ℓ_a + ℓ_b, x_μ, x_rest
 end
 
@@ -297,7 +297,7 @@ end
 function transport_to_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector) where {S<:StdMeasure}
     z_a, _, x2 = transport_to_std_with_rest(S, μ.α, x)
     z_b, _, x_rest = transport_to_std_with_rest(S, μ.β, x2)
-    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest))
     return vcat(z_a, z_b), x_μ, x_rest
 end
 
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 74f23c56..1bc9ba4b 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -70,7 +70,7 @@ function batched_rand_impl(ctx::GenContext, μ::PowerMeasure, sz::Dims)
     _pwr_batched_rand(ctx, μ, sz, fixed_stream_size(pwr_base(μ)))
 end
 function _pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::Dims, ::True)
-    batched_rand_impl(ctx, pwr_base(μ), (_dynamic_dims(pwr_size(μ))..., sz...))
+    batched_rand_impl(ctx, pwr_base(μ), (asnonstatic(pwr_size(μ))..., sz...))
 end
 _pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::Dims, ::False) = _batched_rand_pointwise(ctx, μ, sz)
 
@@ -79,10 +79,7 @@ marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes)
 @inline mspace_elsize(μ::PowerMeasure) = pwr_size(μ)
 @inline mspace_flatsize(μ::PowerMeasure) = _cat_sizes(mspace_flatsize(pwr_base(μ)), pwr_size(μ))
 @inline function mspace_flatsize(::Type{<:PowerMeasure{M,A}}) where {M,A<:Tuple{Vararg{StaticOneToLike}}}
-    _cat_sizes(mspace_flatsize(M), _static_axes_size(A))
-end
-@generated function _static_axes_size(::Type{A}) where {A<:Tuple{Vararg{StaticOneToLike}}}
-    :(StaticArrays.Size($(map(T -> T.parameters[1], A.parameters)...)))
+    _cat_sizes(mspace_flatsize(M), axes2size(A))
 end
 
 function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{AbstractArray,N}}) where {N}
@@ -110,15 +107,15 @@ end
 @inline function mspace_ndims(::Type{<:PowerMeasure{M,A}}) where {M,A<:Tuple}
     _pwr_ndims(mspace_ndims(M), fieldcount(A), fixed_stream_size(M))
 end
-@inline _pwr_ndims(n::Integer, k::Integer, ::Any) = n + k
-@inline _pwr_ndims(::NoMSpaceElementSize, ::Integer, ::True) = 1
-@inline _pwr_ndims(n::NoMSpaceElementSize, ::Integer, ::False) = n
+@inline _pwr_ndims(n::IntegerLike, k::IntegerLike, ::Any) = n + k
+@inline _pwr_ndims(::NoMSpaceElementSize, ::IntegerLike, ::True) = static(1)
+@inline _pwr_ndims(n::NoMSpaceElementSize, ::IntegerLike, ::False) = n
 
 # Local measures of powers at nested variates are products of the local
 # measures of the elements:
 @inline localmeasure(μ::PowerMeasure, ::AbstractArray{<:Number}) = μ
 function localmeasure(μ::PowerMeasure, x::AbstractArray)
-    size(x) == _dynamic_dims(pwr_size(μ)) || return μ
+    size(x) == asnonstatic(pwr_size(μ)) || return μ
     productmeasure(map(Base.Fix1(localmeasure, pwr_base(μ)), x))
 end
 @inline fixed_stream_size(::Type{<:PowerMeasure{M}}) where {M} = fixed_stream_size(M)
@@ -138,7 +135,7 @@ end
     _powered_kernel_impl(f, μ, X, _static_ndims(pwr_base(μ)))
 end
 @inline function _powered_kernel_impl(f::F, μ::PowerMeasure, X, ::Any) where {F}
-    _sum_leading_dims(_batched_kernel(f, pwr_base(μ), X), static(length(pwr_axes(μ))))
+    sum_leading_dims(_batched_kernel(f, pwr_base(μ), X), static(length(pwr_axes(μ))))
 end
 # Numeric batches of powers of bases without a variate rank are batches of
 # streams:
@@ -154,7 +151,7 @@ _powered_stream_kernel(μ::PowerMeasure, X::AbstractArray, ::False) = _streamwis
 
 # Flat batches of powers have the power dimensions after the variate
 # dimensions of the base measure (where the rank of the base is known):
-@inline _check_pwr_batch(X::AbstractArray, μ::PowerMeasure) = _check_pwr_dims(X, _static_ndims(pwr_base(μ)), _dynamic_dims(pwr_size(μ)), false)
+@inline _check_pwr_batch(X::AbstractArray, μ::PowerMeasure) = _check_pwr_dims(X, _static_ndims(pwr_base(μ)), asnonstatic(pwr_size(μ)), false)
 @inline _check_pwr_batch(::Any, ::PowerMeasure) = nothing
 @inline function _check_pwr_dims(X::AbstractArray, ::StaticInteger{K}, dims::Dims, exact::Bool) where {K}
     n = length(dims)
@@ -164,7 +161,6 @@ _powered_stream_kernel(μ::PowerMeasure, X::AbstractArray, ::False) = _streamwis
     return nothing
 end
 @inline _check_pwr_dims(::AbstractArray, ::NoMSpaceElementSize, ::Dims, ::Bool) = nothing
-@inline _dynamic_dims(sz::SizeLike) = map(dynamic, _size_dims(sz))
 @inline batched_logdensityof_impl(μ::PowerMeasure, X) = _powered_kernel(logdensityof_impl, μ, X)
 @inline batched_logdensity_def(μ::PowerMeasure, X) = _powered_kernel(logdensity_def, μ, X)
 
@@ -210,8 +206,8 @@ function batched_logdensityof_with_rest(μ::PowerMeasure, x::AbstractVector, sz:
     _powered_ld_with_rest(μ, x, sz, fixed_stream_size(pwr_base(μ)))
 end
 function _powered_ld_with_rest(μ::PowerMeasure, X::AbstractArray, sz::Dims, ::True)
-    ℓ, X_rest = batched_logdensityof_with_rest(pwr_base(μ), X, (_dynamic_dims(pwr_size(μ))..., sz...))
-    return _sum_leading_dims(ℓ, static(length(pwr_axes(μ)))), X_rest
+    ℓ, X_rest = batched_logdensityof_with_rest(pwr_base(μ), X, (asnonstatic(pwr_size(μ))..., sz...))
+    return sum_leading_dims(ℓ, static(length(pwr_axes(μ)))), X_rest
 end
 function _powered_ld_with_rest(μ::PowerMeasure, x::AbstractVector, ::Tuple{}, ::False)
     ν = pwr_base(μ)
@@ -274,7 +270,7 @@ end
 
 @inline function _check_pwr_variate(μ::PowerMeasure, x::AbstractArray)
     if maybestatic_size(x) != pwr_size(μ)
-        _check_pwr_flat(x, _static_ndims(pwr_base(μ)), _dynamic_dims(pwr_size(μ)))
+        _check_pwr_flat(x, _static_ndims(pwr_base(μ)), asnonstatic(pwr_size(μ)))
     end
     return nothing
 end
@@ -300,7 +296,7 @@ function batched_transport_to_std(::Type{S}, μ::PowerMeasure, X::Union{Tuple,Na
 end
 @inline function _pwr_batched_to_std(::Type{S}, μ::PowerMeasure, X, ::Any) where {S}
     ν, n = _pwr_unwrap(μ)
-    _merge_leading_dims(batched_transport_to_std(S, ν, X), static(1) + n)
+    merge_leading_dims(batched_transport_to_std(S, ν, X), static(1) + n)
 end
 # Numeric batches of powers of bases without a variate rank are batches of
 # streams:
@@ -316,14 +312,14 @@ function batched_transport_from_std(::Type{S}, μ::PowerMeasure, Z::AbstractArra
     n_rows = _batch_dims(Z)[1]
     dof_ν = _base_dof(n_rows, prod(dims))
     dof_ν * prod(dims) == n_rows || _throw_std_length_mismatch()
-    batched_transport_from_std(S, ν, _reshape_batch(Z, (dof_ν, dims..., Base.tail(_batch_dims(Z))...)))
+    batched_transport_from_std(S, ν, maybestatic_reshape(Z, (dof_ν, dims..., Base.tail(_batch_dims(Z))...)))
 end
 
 # Empty powers leave the degrees of freedom of the base undetermined:
 @inline _base_dof(n_rows::IntegerLike, n_pwr::IntegerLike) = n_rows ÷ max(n_pwr, one(n_pwr))
 
 # All power dimensions of nested powers, innermost first:
-@inline _pwr_dims(μ::PowerMeasure) = (_pwr_dims(pwr_base(μ))..., _size_dims(pwr_size(μ))...)
+@inline _pwr_dims(μ::PowerMeasure) = (_pwr_dims(pwr_base(μ))..., size_dims(pwr_size(μ))...)
 @inline _pwr_dims(ν) = ()
 
 # Point transport: flat variates are batches with zero batch dimensions,
@@ -350,7 +346,7 @@ function batched_transport_to_std_with_rest(::Type{S}, μ::PowerMeasure, X::Abst
     _pwr_to_std_with_rest(S, μ, X, sz, fixed_stream_size(pwr_base(μ)))
 end
 @inline function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, X::AbstractArray, sz::Dims, ::True) where {S}
-    batched_transport_to_std_with_rest(S, pwr_base(μ), X, (_dynamic_dims(pwr_size(μ))..., sz...))
+    batched_transport_to_std_with_rest(S, pwr_base(μ), X, (asnonstatic(pwr_size(μ))..., sz...))
 end
 function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::Tuple{}, ::False) where {S}
     z, _, x_rest = transport_to_std_with_rest(S, μ, x)
@@ -374,7 +370,7 @@ end
 # batched protocol:
 function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::NoMSpaceElementSize) where {S}
     z, x_rest = _pwr_to_std_with_rest(S, μ, x, (), static(true))
-    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest))
     return z, x_μ, x_rest
 end
 function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::False) where {S}
@@ -384,7 +380,7 @@ function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVec
     for i in eachindex(zs)
         zs[i], _, x_rest = transport_to_std_with_rest(S, ν, x_rest)
     end
-    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest))
     return reduce(vcat, [z for z in zs]), x_μ, x_rest
 end
 
@@ -400,7 +396,7 @@ end
 
 # The stream length of a power with a base of fixed stream length:
 @inline function _fixed_stream_length(μ::PowerMeasure)
-    _fixed_stream_length(pwr_base(μ)) * prod(_dynamic_dims(pwr_size(μ)))
+    _fixed_stream_length(pwr_base(μ)) * prod(asnonstatic(pwr_size(μ)))
 end
 
 # The nested variate layout of a power over its flat storage, batches of
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index f90d065d..ae8b05d0 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -237,7 +237,7 @@ end
 @inline function _array_product_kernel(f::F, μ::ProductMeasure, X::AbstractArray, ::StaticInteger{0}) where {F}
     mar = marginals(μ)
     _check_flatsize(X, maybestatic_size(mar))
-    _sum_leading_dims(_marginal_broadcast(_DynamicPointLogd(f), mar, X), static(ndims(mar)))
+    sum_leading_dims(_marginal_broadcast(_DynamicPointLogd(f), mar, X), static(ndims(mar)))
 end
 @inline function _array_product_kernel(f::F, μ::ProductMeasure, X::AbstractArray, ::StaticInteger{K}) where {F,K}
     _marginal_slices_ld(f, marginals(μ), X, Val(K), Val(ndims(X) - K - ndims(marginals(μ))))
@@ -275,7 +275,7 @@ end
 @inline function logdensity_def(μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray{<:Number}) where {M}
     _array_product_ld(logdensity_def, μ, x, mspace_ndims(M))
 end
-@inline function _array_product_ld(f::F, μ::ProductMeasure, x::AbstractArray, ::Integer) where {F}
+@inline function _array_product_ld(f::F, μ::ProductMeasure, x::AbstractArray, ::IntegerLike) where {F}
     _point_result(_materialize(_batched_kernel(f, μ, x)), μ)
 end
 @inline function _array_product_ld(f::F, μ::ProductMeasure, x::AbstractArray, ::NoMSpaceElementSize) where {F}
@@ -346,7 +346,7 @@ fast_dof(d::AbstractProductMeasure) = _sum_dofs(fast_dof, marginals(d))
 # of freedom each, so their total needs no reduction over the marginals
 # (which may live on a device):
 @inline function _unit_dof(::Type{M}) where {M}
-    static(mspace_ndims(M) === 0 && preferred_stdmeasure(M) isa Type{<:StdMeasure})
+    static(_static_ndims_of(mspace_ndims(M)) === static(0) && preferred_stdmeasure(M) isa Type{<:StdMeasure})
 end
 @inline _sum_dofs(f, mar::StaticArray) = mapreduce(f, +, mar; init = static(0))
 @inline _dynamic_dof(n::IntegerLike) = dynamic(n)
@@ -420,7 +420,7 @@ end
 # Streams of tuple product variates are consumed marginal by marginal:
 function transport_to_std_with_rest(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, x::AbstractVector) where {S<:StdMeasure}
     z, x_rest = _marginals_to_std_with_rest(S, values(marginals(μ)), x)
-    x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest))
+    x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest))
     return z, x_μ, x_rest
 end
 
@@ -556,7 +556,7 @@ end
 function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{true}, ::Any) where {S}
     mar = marginals(μ)
     size(Z, 1) == length(mar) || _throw_std_length_mismatch()
-    _materialize(_marginal_broadcast(_FromStd{S}(), mar, _reshape_batch(Z, (_batch_dims(mar)..., Base.tail(_batch_dims(Z))...))))
+    _materialize(_marginal_broadcast(_FromStd{S}(), mar, maybestatic_reshape(Z, (_batch_dims(mar)..., Base.tail(_batch_dims(Z))...))))
 end
 function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{false}, ::StaticInteger{K}) where {S,K}
     X, Z_rest = _marginals_from_std_loop(S, marginals(μ), Z, Val(K))
@@ -689,13 +689,7 @@ function batched_logdensityof_with_rest(μ::ProductMeasure{<:NamedTuple{names}},
 end
 
 # Folded pairwise over the marginal types, so that the result is a constant:
-@inline fixed_stream_size(::Type{<:ProductMeasure{M}}) where {M<:Tuple} = _all_fixed_stream_sizes(M)
-@inline _all_fixed_stream_sizes(::Type{Tuple{}}) = static(true)
-@inline function _all_fixed_stream_sizes(::Type{M}) where {M<:Tuple}
-    _both(fixed_stream_size(Base.tuple_type_head(M)), _all_fixed_stream_sizes(Base.tuple_type_tail(M)))
-end
-@inline _both(::True, ::True) = static(true)
-@inline _both(::StaticBool, ::StaticBool) = static(false)
+@inline fixed_stream_size(::Type{<:ProductMeasure{M}}) where {M<:Tuple} = static_all(fixed_stream_size, M)
 @inline function fixed_stream_size(::Type{<:ProductMeasure{NamedTuple{names,M}}}) where {names,M<:Tuple}
     fixed_stream_size(ProductMeasure{M})
 end
diff --git a/src/combinators/reshape.jl b/src/combinators/reshape.jl
index ad83e472..e20a0c5f 100644
--- a/src/combinators/reshape.jl
+++ b/src/combinators/reshape.jl
@@ -63,7 +63,7 @@ mreshape(m::AbstractMeasure, sz::SizeLike) = pushfwd(Reshape(sz, some_mspace_els
     _reshaped_flatsize(mspace_flatsize(μ.origin), mspace_elsize(μ.origin), μ.f.output_size)
 end
 @inline function _reshaped_flatsize(sz_flat::SizeLike, sz_outer::SizeLike, sz_out)
-    _reshaped_flatsize(Val(length(_size_dims(sz_flat)) == length(_size_dims(sz_outer))), sz_out)
+    _reshaped_flatsize(Val(length(size_dims(sz_flat)) == length(size_dims(sz_outer))), sz_out)
 end
 @inline _reshaped_flatsize(::Val{true}, sz_out) = sz_out
 @inline _reshaped_flatsize(::Val{false}, sz_out) = NoMSpaceElementSize{typeof(sz_out)}()
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index db7a8950..53e6fde7 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -173,24 +173,18 @@ end
 
 @inline mspace_flatsize(μ::SuperpositionMeasure) = mspace_flatsize(typeof(μ))
 
-# The variate rank of a superposition is the common rank of its components:
+# The variate rank of a superposition is the common rank of its components,
+# folded pairwise over the component types so that it stays a constant:
 @inline mspace_ndims(::Type{<:SuperpositionMeasure{C}}) where {C<:AbstractArray} = mspace_ndims(eltype(C))
-@generated function mspace_ndims(::Type{MU}) where {C<:Tuple,MU<:SuperpositionMeasure{C}}
-    args = [:(mspace_ndims($T)) for T in C.parameters]
-    :(_common_ndims(($(args...),), MU))
-end
-# Pairwise comparisons fold to a constant rank where `all` doesn't (Julia 1.10):
-@inline _common_ndims(ns::Tuple{Integer,Vararg{Integer}}, ::Type{MU}) where {MU} = _common_ndims_of(first(ns), Base.tail(ns), MU)
-@inline _common_ndims(::Tuple, ::Type{MU}) where {MU} = NoMSpaceElementSize{MU}()
-@inline _common_ndims_of(n::Integer, ::Tuple{}, ::Type) = n
-@inline function _common_ndims_of(n::Integer, ns::Tuple{Integer,Vararg{Integer}}, ::Type{MU}) where {MU}
-    n == first(ns) ? _common_ndims_of(n, Base.tail(ns), MU) : NoMSpaceElementSize{MU}()
+@inline function mspace_ndims(::Type{MU}) where {C<:Tuple,MU<:SuperpositionMeasure{C}}
+    static_reduce(_CommonNDims{MU}(), mspace_ndims, C)
 end
+struct _CommonNDims{MU} <: Function end
+@inline (::_CommonNDims{MU})(a::IntegerLike, b::IntegerLike) where {MU} = a == b ? a : NoMSpaceElementSize{MU}()
+@inline (::_CommonNDims{MU})(::Any, ::Any) where {MU} = NoMSpaceElementSize{MU}()
+
 @inline mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:AbstractArray} = _scalar_or_unknown(mspace_flatsize(eltype(C)))
-@inline mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:Tuple} = _common_scalar_flatsize(C)
-@generated function _common_scalar_flatsize(::Type{C}) where {C<:Tuple}
-    args = [:(mspace_flatsize($T)) for T in C.parameters]
-    :(_all_scalar_sizes($(args...)))
+@inline function mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:Tuple}
+    _scalar_or_unknown(static_reduce(_common_flatsize, mspace_flatsize, C))
 end
-@inline _all_scalar_sizes(::Tuple{}...) = ()
-@inline _all_scalar_sizes(szs...) = NoMSpaceElementSize{typeof(szs)}()
+@inline _common_flatsize(a, b) = a === b ? a : NoMSpaceElementSize{typeof((a, b))}()
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index b0faffa5..1b5d7239 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -206,7 +206,7 @@ end
 function _elementwise_pushfwd_ld(f::F, ν::PushforwardMeasure, Y, k::StaticInteger) where {F}
     f_inv = ν.finv.f
     ℓ = _batched_kernel(f, ν.origin, broadcast(f_inv, Y))
-    ladj = _sum_leading_dims(broadcast(_LadjOf(f_inv), Y), k)
+    ladj = sum_leading_dims(broadcast(_LadjOf(f_inv), Y), k)
     return _lazy_combine_ladj(ℓ, ladj)
 end
 function _elementwise_pushfwd_ld(f::F, ν::PushforwardMeasure, Y, ::NoMSpaceElementSize) where {F}
diff --git a/src/density-batched.jl b/src/density-batched.jl
index e5748699..61e9f6a0 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -61,9 +61,9 @@ end
 
 # The variate rank as a static integer, from the type where known:
 @inline _static_ndims(μ::MU) where {MU} = _static_ndims(mspace_ndims(MU), μ)
-@inline _static_ndims(n::Integer, μ) = static(n)
+@inline _static_ndims(n::IntegerLike, μ) = static(n)
 @inline _static_ndims(::NoMSpaceElementSize, μ) = _static_ndims_of(mspace_ndims(μ))
-@inline _static_ndims_of(n::Integer) = static(n)
+@inline _static_ndims_of(n::IntegerLike) = static(n)
 @inline _static_ndims_of(n::NoMSpaceElementSize) = n
 
 """
@@ -79,7 +79,7 @@ function batched_logdensity_def end
     _default_batched_kernel(logdensity_def, μ, X, _static_ndims(μ))
 end
 
-@inline _default_batched_kernel(f::F, μ, X, n::Integer) where {F} = _default_batched_kernel(f, μ, X, static(n))
+@inline _default_batched_kernel(f::F, μ, X, n::IntegerLike) where {F} = _default_batched_kernel(f, μ, X, static(n))
 @inline _default_batched_kernel(f::F, μ, X, ::StaticInteger{0}) where {F} = _scalar_kernel_broadcast(f, μ, X)
 @inline _default_batched_kernel(f::F, μ, X::AbstractArray, ::StaticInteger{0}) where {F} = _scalar_kernel_broadcast(f, μ, X)
 @inline function _default_batched_kernel(f::F, μ, X::AbstractArray, ::StaticInteger{K}) where {F,K}
@@ -199,8 +199,8 @@ const _LazyBroadcast = Broadcast.Broadcasted
 
 # The leading dimensions of a flat batch must match a flat variate size:
 @inline function _check_flatsize(A::AbstractArray, sz_flat::SizeLike)
-    n = length(_size_dims(sz_flat))
-    if ndims(A) < n || ntuple(i -> size(A, i), Val(n)) != Tuple(_size_dims(sz_flat))
+    n = length(size_dims(sz_flat))
+    if ndims(A) < n || ntuple(i -> size(A, i), Val(n)) != asnonstatic(sz_flat)
         _throw_size_mismatch()
     end
     return nothing
@@ -210,51 +210,12 @@ end
 @inline _materialize(x) = x
 
 
-# Lazy sums over the leading `N` dimensions; a full reduction yields a
-# number. Lazy broadcasts are reduced without materialization where the
-# broadcast style supports it, and materialized before partial reductions.
-const _EagerReducibleBroadcast = Broadcast.Broadcasted{<:Union{Broadcast.DefaultArrayStyle,StaticArrays.StaticArrayStyle}}
-
-@inline _sum_leading_dims(x::Number, ::StaticInteger{0}) = x
-@noinline function _sum_leading_dims(::Number, ::StaticInteger)
-    throw(ArgumentError("Variates of powers of measures must be arrays"))
-end
-@inline _sum_leading_dims(A::AbstractArray, n::StaticInteger) = _sum_leading_dims_impl(A, n, static(ndims(A)))
-@inline _sum_leading_dims(bc::_LazyBroadcast, n::StaticInteger) = _sum_leading_dims_lazy(bc, n, static(ndims(bc)))
-@inline _sum_leading_dims_lazy(bc::_LazyBroadcast, ::StaticInteger{0}, ::StaticInteger) = bc
-@inline _sum_leading_dims_lazy(bc::_LazyBroadcast, ::StaticInteger{0}, ::StaticInteger{0}) = bc
-@inline _sum_leading_dims_lazy(bc::_EagerReducibleBroadcast, ::StaticInteger{0}, ::StaticInteger{0}) = bc
-@inline function _sum_leading_dims_lazy(bc::_EagerReducibleBroadcast, ::StaticInteger{N}, ::StaticInteger{N}) where {N}
-    length(bc) == 0 ? sum(copy(bc)) : sum(bc)
-end
-@inline _sum_leading_dims_lazy(bc::_LazyBroadcast, ::StaticInteger{N}, ::StaticInteger{N}) where {N} = sum(copy(bc))
-@inline function _sum_leading_dims_lazy(bc::_LazyBroadcast, n::StaticInteger, ::StaticInteger)
-    _sum_leading_dims(copy(bc), n)
-end
-@inline _sum_leading_dims_impl(A::AbstractArray, ::StaticInteger{0}, ::StaticInteger) = A
-@inline _sum_leading_dims_impl(A::AbstractArray, ::StaticInteger{0}, ::StaticInteger{0}) = A
-@inline _sum_leading_dims_impl(A::AbstractArray, ::StaticInteger{N}, ::StaticInteger{N}) where {N} = sum(A)
-@inline function _sum_leading_dims_impl(A::AbstractArray, ::StaticInteger{N}, ::StaticInteger) where {N}
-    _drop_leading_dims(_sum_dims_seq(A, static(N)), static(N))
-end
-
-# Drops the leading `N` (singleton) dimensions by reshaping, which keeps
-# static arrays static and infers where `dropdims` doesn't:
-@inline function _drop_leading_dims(A::AbstractArray, ::StaticInteger{N}) where {N}
-    dims = _batch_dims(A)
-    _reshape_batch(A, ntuple(i -> dims[N + i], Val(length(dims) - N)))
-end
-@inline _sum_dims_seq(A::AbstractArray, ::StaticInteger{0}) = A
-@inline function _sum_dims_seq(A::AbstractArray, ::StaticInteger{N}) where {N}
-    _sum_dims_seq(sum(A; dims = N), static(N - 1))
-end
-
 @inline _lazy_add(a::Number, b::Number) = a + b
 @inline _lazy_add(a, b) = Broadcast.instantiate(Broadcast.broadcasted(+, a, b))
 
 # Zero log-densities over the batch dimensions of a flat batch of variates
 # with `n` variate dimensions:
-@inline function _zero_logd_batch(X::AbstractArray, n::Integer)
+@inline function _zero_logd_batch(X::AbstractArray, n::IntegerLike)
     FillArrays.Zeros{_logd_numtype(X)}(ntuple(i -> size(X, n + i), ndims(X) - n))
 end
 @inline _zero_logd_batch(X::AbstractArray{<:Any,N}, ::StaticInteger{N}) where {N} = zero(_logd_numtype(X))
@@ -320,23 +281,21 @@ end
     return _reshape_consumed(X_flat, (dims..., sz...)), X_rest
 end
 @inline _consumed_dims(::Tuple{}) = (static(1),)
-@inline _consumed_dims(vsz::SizeLike) = _size_dims(vsz)
+@inline _consumed_dims(vsz::SizeLike) = size_dims(vsz)
 @inline _chunk_rows(n::IntegerLike, ::Tuple{}) = n
 @inline _chunk_rows(n::IntegerLike, sz::Dims) = dynamic(n) * prod(sz)
 
 @inline _reshape_consumed(X_flat::AbstractArray, ::Tuple{IntegerLike}) = X_flat
 @inline function _reshape_consumed(X_flat::AbstractArray, dims::Tuple{Vararg{IntegerLike}})
-    _reshape_batch(X_flat, (dims..., Base.tail(_batch_dims(X_flat))...))
+    maybestatic_reshape(X_flat, (dims..., Base.tail(_batch_dims(X_flat))...))
 end
 
-# Sizes as tuples of (maybe static) integers, reshapes that keep static
-# arrays static, and the leading dimension of a batch of streams:
-@inline _batch_dims(A::AbstractArray) = _size_dims(maybestatic_size(A))
-@inline _reshape_batch(A::AbstractArray, dims::Tuple) = reshape(A, map(dynamic, dims))
-@inline _reshape_batch(A::StaticArray, dims::Tuple{Vararg{StaticInteger}}) = maybestatic_reshape(A, dims)
-@inline _as_stdstream_batch(Z::AbstractArray) = _reshape_batch(Z, (static(1), _batch_dims(Z)...))
+# Sizes as tuples of (maybe static) integers and the leading dimension of
+# a batch of streams:
+@inline _batch_dims(A::AbstractArray) = size_dims(maybestatic_size(A))
+@inline _as_stdstream_batch(Z::AbstractArray) = merge_leading_dims(Z, static(0))
 @inline _as_stdstream_batch(z::Number) = SVector(z)
-@inline _drop_stdstream_dim(Z::AbstractArray) = _reshape_batch(Z, Base.tail(_batch_dims(Z)))
+@inline _drop_stdstream_dim(Z::AbstractArray) = drop_leading_dims(Z, static(1))
 
 @inline function _batched_split(A::AbstractArray, n::IntegerLike)
     n_rows = dynamic(n)
@@ -352,11 +311,7 @@ end
 end
 
 # Static streams split into static chunks for static row counts:
-@inline function _batched_split(A::StaticVector, n_rows::StaticInteger{N}) where {N}
-    idxs = maybestatic_eachindex(A)
-    i0 = maybestatic_first(idxs)
-    _get_or_view(A, i0, i0 + n_rows - static(1)), _get_or_view(A, i0 + n_rows, maybestatic_last(idxs))
-end
+@inline _batched_split(A::StaticVector, n_rows::StaticInteger) = split_at(A, n_rows)
 
 @noinline function _throw_stream_too_long()
     throw(ArgumentError("Variate streams too long during density evaluation"))
@@ -367,7 +322,7 @@ end
 # in fused operations; otherwise a batch of streams is consumed stream by
 # stream by the outermost stream combinator.
 @inline fixed_stream_size(μ::MU) where {MU} = fixed_stream_size(MU)
-@inline fixed_stream_size(::Type{MU}) where {MU} = static(mspace_ndims(MU) isa Integer)
+@inline fixed_stream_size(::Type{MU}) where {MU} = static(mspace_ndims(MU) isa IntegerLike)
 
 # Batches of streams consumed stream by stream (host loop):
 function _streamwise_ld(f::F, μ, X::AbstractArray) where {F}
diff --git a/src/mspace.jl b/src/mspace.jl
index 4ddb28ef..01a6f7a8 100644
--- a/src/mspace.jl
+++ b/src/mspace.jl
@@ -45,9 +45,7 @@ function mspace_flatsize end
 
 @inline mspace_flatsize(μ::AbstractMeasure) = NoMSpaceElementSize{typeof(μ)}()
 
-@inline _cat_sizes(a::SizeLike, b::SizeLike) = canonical_size((_size_dims(a)..., _size_dims(b)...))
-@inline _size_dims(sz::Tuple) = sz
-@inline _size_dims(::StaticArrays.Size{S}) where {S} = map(static, S)
+@inline _cat_sizes(a::SizeLike, b::SizeLike) = canonical_size((size_dims(a)..., size_dims(b)...))
 @inline _cat_sizes(a::NoMSpaceElementSize, ::SizeLike) = a
 @inline _cat_sizes(::SizeLike, b::NoMSpaceElementSize) = b
 @inline _cat_sizes(a::NoMSpaceElementSize, ::NoMSpaceElementSize) = a
@@ -120,10 +118,10 @@ function mspace_ndims end
 @inline mspace_ndims(::Type{MU}) where {MU} = _ndims_of_size(mspace_flatsize(MU), MU)
 @inline mspace_ndims(μ::MU) where {MU} = _ndims_of_size(mspace_flatsize(μ), MU, mspace_ndims(MU))
 
-@inline _ndims_of_size(sz::SizeLike, ::Type) = length(_size_dims(sz))
+@inline _ndims_of_size(sz::SizeLike, ::Type) = maybestatic_length(size_dims(sz))
 @inline _ndims_of_size(::NoMSpaceElementSize, ::Type{MU}) where {MU} = NoMSpaceElementSize{MU}()
-@inline _ndims_of_size(sz::SizeLike, ::Type, ::Any) = length(_size_dims(sz))
+@inline _ndims_of_size(sz::SizeLike, ::Type, ::Any) = maybestatic_length(size_dims(sz))
 @inline _ndims_of_size(::NoMSpaceElementSize, ::Type, n) = n
 
-@inline _add_ndims(n::Integer, k::Integer) = n + k
-@inline _add_ndims(n::NoMSpaceElementSize, ::Integer) = n
+@inline _add_ndims(n::IntegerLike, k::IntegerLike) = n + k
+@inline _add_ndims(n::NoMSpaceElementSize, ::IntegerLike) = n
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index f5025bb9..4fdbb610 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -83,16 +83,11 @@ end
 
 # Batches of array variates: all elements of a variate must match.
 function batched_logdensityof_impl(μ::Dirac{<:AbstractArray{<:Number,N}}, X::AbstractArray) where {N}
-    matches = _all_leading_dims(X .== μ.x, static(N))
+    matches = all_leading_dims(X .== μ.x, static(N))
     ifelse.(matches, zero(_logd_numtype(X)), _neg_inf_logd(X))
 end
 function batched_logdensity_def(μ::Dirac{<:AbstractArray{<:Number}}, X::AbstractArray)
     _zero_logd_batch(X, static(ndims(μ.x)))
 end
 
-@inline _all_leading_dims(A::AbstractArray{Bool,N}, ::StaticInteger{N}) where {N} = all(A)
-@inline function _all_leading_dims(A::AbstractArray{Bool}, ::StaticInteger{N}) where {N}
-    _drop_leading_dims(all(A; dims = ntuple(identity, Val(N))), static(N))
-end
-
 Adapt.adapt_structure(to, μ::Dirac) = Dirac(Adapt.adapt(to, μ.x))
diff --git a/src/standard/stdconvert.jl b/src/standard/stdconvert.jl
index 520f3097..8ddf841d 100644
--- a/src/standard/stdconvert.jl
+++ b/src/standard/stdconvert.jl
@@ -59,5 +59,5 @@ end
 function batched_transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, X::AbstractArray) where {NU<:StdMeasure,MU<:StdMeasure}
     n_μ = length(mspace_flatsize(μ))
     batch_dims = ntuple(i -> size(X, n_μ + i), Val(ndims(X) - n_μ))
-    reshape(stdconvert(NU, MU, X), (map(dynamic, _size_dims(mspace_flatsize(ν)))..., batch_dims...))
+    reshape(stdconvert(NU, MU, X), (asnonstatic(mspace_flatsize(ν))..., batch_dims...))
 end
diff --git a/src/transport-batched.jl b/src/transport-batched.jl
index 4b476bfb..79a1e993 100644
--- a/src/transport-batched.jl
+++ b/src/transport-batched.jl
@@ -45,22 +45,15 @@ end
     stacked(map(Base.Fix1(_ToStd{S}(), μ), sliced(X, Val(K))))
 end
 
-# Merge the leading `N` dimensions of an array into one, `N == 0` adds a
-# leading dimension of size one:
-@inline function _merge_leading_dims(A::AbstractArray, ::StaticInteger{N}) where {N}
-    ndims(A) >= N || _throw_size_mismatch()
-    dims = _batch_dims(A)
-    lead = ntuple(i -> dims[i], Val(N))
-    _reshape_batch(A, (prod(lead), ntuple(i -> dims[N + i], Val(length(dims) - N))...))
-end
-@inline _merge_leading_dims(A::AbstractArray, ::StaticInteger{0}) = _reshape_batch(A, (static(1), _batch_dims(A)...))
-
 # A flat batch of variates of `μ` as a batch of streams, the variate
 # dimensions merged into the first dimension. Tuples of batches (tuple
 # products and their powers) interleave the rows of their components
 # variate by variate.
 @inline _as_stream_batch(X, μ) = _as_stream_batch(X, _static_ndims(μ))
-@inline _as_stream_batch(X::AbstractArray, ::StaticInteger{K}) where {K} = _merge_leading_dims(X, static(K))
+@inline function _as_stream_batch(X::AbstractArray, ::StaticInteger{K}) where {K}
+    ndims(X) >= K || _throw_size_mismatch()
+    merge_leading_dims(X, static(K))
+end
 @inline _as_stream_batch(x::Number, ::StaticInteger{0}) = SVector(x)
 @noinline function _as_stream_batch(X, ::NoMSpaceElementSize)
     throw(ArgumentError("Concatenating batches of variates requires MeasureBase.mspace_ndims to be declared for the measures involved"))
@@ -71,12 +64,12 @@ end
 function _as_stream_batch(X::Union{Tuple,NamedTuple}, μ::PowerMeasure)
     ν, _ = _pwr_unwrap(μ)
     n_pwr = length(_pwr_dims(μ))
-    n = prod(map(dynamic, _pwr_dims(μ)))
+    n = prod(asnonstatic(_pwr_dims(μ)))
     parts = map(values(X), values(marginals(ν))) do Xi, m
         A = _as_stream_batch(Xi, m)
         reshape(A, (size(A, 1), n, ntuple(i -> size(A, 1 + n_pwr + i), Val(ndims(A) - 1 - n_pwr))...))
     end
-    _merge_leading_dims(vcat(parts...), static(2))
+    merge_leading_dims(vcat(parts...), static(2))
 end
 
 # The standard variates of a single variate must form a vector:
@@ -163,10 +156,10 @@ end
 # Standard variates of `prod(sz)` variates per stream, `(dof, sz..., batch
 # dims...)`, as one stream chunk `(dof * prod(sz), batch dims...)`, and
 # back:
-@inline _merge_multiplicity(Z::AbstractArray, sz::Dims) = _merge_leading_dims(Z, static(1) + static(length(sz)))
+@inline _merge_multiplicity(Z::AbstractArray, sz::Dims) = merge_leading_dims(Z, static(1) + static(length(sz)))
 @inline _split_multiplicity(Z::AbstractArray, ::Tuple{}, n) = Z
 @inline function _split_multiplicity(Z::AbstractArray, sz::Dims, n)
-    _reshape_batch(Z, (n, sz..., Base.tail(_batch_dims(Z))...))
+    maybestatic_reshape(Z, (n, sz..., Base.tail(_batch_dims(Z))...))
 end
 
 
diff --git a/src/transport.jl b/src/transport.jl
index 4c7403c0..6061a21c 100644
--- a/src/transport.jl
+++ b/src/transport.jl
@@ -289,7 +289,7 @@ function _from_std_with_rest_bydof(::Type{S}, μ, z::AbstractVector, n::IntegerL
     if maybestatic_length(z) < n
         throw(ArgumentError("Stream of standard variates too short during transport"))
     end
-    z_μ, z_rest = _split_after(z, n)
+    z_μ, z_rest = split_at(z, n)
     return transport_from_std(S, μ, _chunk_as_variate(μ, z_μ)), z_rest
 end
 

From c06f1cbc2fb027163fba5b79d0a38bdf9a225929 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 21 Sep 2026 16:25:49 +0200
Subject: [PATCH 115/122] Keep static variate sizes static end to end

The multiplicity of the with-rest protocol and the batch size of random
variate generation were `Dims`, which turned every statically sized
power into a dynamic one as soon as it was consumed from a stream. Both
are `SizeLike` now and stay static: `_chunk_rows` multiplies with
`size2length`, powers hand their `pwr_size` to their base unchanged,
combined measures and tuple products split and reshape their rows with
`_chunk_rows` and `maybestatic_reshape`, and `_fixed_stream_length` keeps
static lengths. The static split of static vectors is therefore reachable
through with-rest.

Bulk draws of a fully static batch size return static arrays on the CPU,
so `rand` of a statically sized power, product or bind allocates nothing.
Other compute units keep allocating their own arrays.

Standard measures check their variates directly: going through their base
measures makes the inliner cut the recursion, and the resulting call
boxes the variate of every scalar marginal of a tuple product.

Created by generative AI.
---
 .../MeasureBaseDistributionsExt.jl            |  2 +-
 .../distribution_measure.jl                   |  8 +--
 src/MeasureBase.jl                            |  4 +-
 src/combinators/bind.jl                       |  8 +--
 src/combinators/combined.jl                   | 31 ++++-----
 src/combinators/half.jl                       |  2 +-
 src/combinators/power.jl                      | 30 ++++-----
 src/combinators/product.jl                    | 30 ++++-----
 src/combinators/spikemixture.jl               |  6 +-
 src/combinators/superpose.jl                  |  6 +-
 src/combinators/transformedmeasure.jl         |  6 +-
 src/combinators/weighted.jl                   |  6 +-
 src/density-batched.jl                        | 11 ++-
 src/primitives/dirac.jl                       |  2 +-
 src/rand.jl                                   | 67 ++++++++++++-------
 src/standard/stdexponential.jl                |  2 +-
 src/standard/stdlogistic.jl                   |  2 +-
 src/standard/stdmeasure.jl                    |  5 ++
 src/standard/stdnormal.jl                     |  2 +-
 src/standard/stduniform.jl                    |  2 +-
 src/transport-batched.jl                      | 20 +++---
 21 files changed, 136 insertions(+), 116 deletions(-)

diff --git a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
index c0ca18f0..3733c025 100644
--- a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
+++ b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl
@@ -49,7 +49,7 @@ using SpecialFunctions: loggamma, logbeta, gamma_inc, gamma_inc_inv, beta_inc, b
 using HeterogeneousComputing: real_numtype, GenContext, get_rng, get_precision, get_compute_unit, CPUnit, AbstractComputeUnit
 
 using Static: True, False, StaticInt, static, dynamic
-using StaticThings: asnonstatic
+using StaticThings: SizeLike, asnonstatic
 using FillArrays: Fill, Ones, Zeros
 
 using ArgCheck: @argcheck
diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
index ddd6aced..acf9fa53 100644
--- a/ext/MeasureBaseDistributionsExt/distribution_measure.jl
+++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl
@@ -18,15 +18,15 @@ const DistributionMeasure{F<:VariateForm,S<:ValueSupport,D<:Distribution{F,S}} =
 # Distributions' samplers run on the CPU, variates on other compute units
 # are generated from standard variates via the transports:
 MeasureBase.rand_impl(ctx::GenContext, m::DistributionMeasure) = _dist_rand(ctx, m, get_compute_unit(ctx))
-MeasureBase.batched_rand_impl(ctx::GenContext, m::DistributionMeasure, sz::Dims) = _dist_batched_rand(ctx, m, sz, get_compute_unit(ctx))
+MeasureBase.batched_rand_impl(ctx::GenContext, m::DistributionMeasure, sz::SizeLike) = _dist_batched_rand(ctx, m, sz, get_compute_unit(ctx))
 
 _dist_rand(ctx::GenContext, m::DistributionMeasure, ::CPUnit) =
     convert_realtype(get_precision(ctx), rand(get_rng(ctx), m.obj))
 _dist_rand(ctx::GenContext, m::DistributionMeasure, ::AbstractComputeUnit) =
     MeasureBase._rand_default(ctx, m, (), MeasureBase._NoRandImpl())
-_dist_batched_rand(ctx::GenContext, m::DistributionMeasure, sz::Dims, ::CPUnit) =
-    _flat_powrand(get_rng(ctx), get_precision(ctx), m.obj, sz)
-_dist_batched_rand(ctx::GenContext, m::DistributionMeasure, sz::Dims, ::AbstractComputeUnit) =
+_dist_batched_rand(ctx::GenContext, m::DistributionMeasure, sz::SizeLike, ::CPUnit) =
+    _flat_powrand(get_rng(ctx), get_precision(ctx), m.obj, asnonstatic(sz))
+_dist_batched_rand(ctx::GenContext, m::DistributionMeasure, sz::SizeLike, ::AbstractComputeUnit) =
     MeasureBase._rand_default(ctx, m, sz, MeasureBase._NoRandImpl())
 
 # A single variate for zero batch dimensions, flat batches otherwise:
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index 976a6731..f584b9f1 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -45,12 +45,12 @@ using FunctionChains
 using PropertyFunctions: PropSelFunction
 
 using StaticThings:
-    SizeLike, OneToLike, StaticOneToLike, IntegerLike,
+    SizeLike, StaticSizeLike, OneToLike, StaticOneToLike, IntegerLike,
     asaxes, asnonstatic, canonical_size, size_dims,
     maybestatic_eachindex, maybestatic_length, maybestatic_size,
     maybestatic_first, maybestatic_last, maybestatic_view,
     maybestatic_oneto, maybestatic_fill, maybestatic_reshape,
-    axes2size, size2length, split_at,
+    axes2size, size2length, split_at, staticarray_type,
     static_all, static_any, static_reduce,
     sum_leading_dims, drop_leading_dims, merge_leading_dims, all_leading_dims
 
diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl
index 39047fc7..2900d999 100644
--- a/src/combinators/bind.jl
+++ b/src/combinators/bind.jl
@@ -349,7 +349,7 @@ batched_logdensityof_impl(μ::_BindBy{typeof(vcat)}, X::AbstractArray) = _stream
 
 # Batches of streams containing binds are consumed stream by stream (by
 # the outermost stream combinator, see `fixed_stream_size`):
-@noinline function batched_logdensityof_with_rest(::Bind, ::AbstractArray, ::Dims)
+@noinline function batched_logdensityof_with_rest(::Bind, ::AbstractArray, ::SizeLike)
     throw(ArgumentError("Batches of variate streams containing binds must be consumed stream by stream"))
 end
 batched_logdensityof_impl(μ::_BindBy{typeof(vcat)}, x::AbstractVector) = _bind_ld_impl(vcat, μ, x)
@@ -363,7 +363,7 @@ end
 
 # The secondary measure depends on the primary variate, so batches are
 # generated variate by variate:
-batched_rand_impl(ctx::GenContext, μ::Bind, sz::Dims) = _batched_rand_pointwise(ctx, μ, sz)
+batched_rand_impl(ctx::GenContext, μ::Bind, sz::SizeLike) = _batched_rand_pointwise(ctx, μ, sz)
 
 
 # Transport consumes the variate parts of the primary and secondary
@@ -412,13 +412,13 @@ end
 # The secondary measure depends on the primary variate, so batches of
 # streams are consumed stream by stream (by the outermost stream
 # combinator, see `fixed_stream_size`):
-function batched_transport_to_std_with_rest(::Type{S}, μ::Bind, X::AbstractArray, sz::Dims) where {S<:StdMeasure}
+function batched_transport_to_std_with_rest(::Type{S}, μ::Bind, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure}
     _bind_to_std_with_rest(S, μ, X, sz)
 end
 function _bind_to_std_with_rest(::Type{S}, μ::Bind, x::AbstractVector, ::Tuple{}) where {S}
     z, _, x_rest = transport_to_std_with_rest(S, μ, x)
     return z, x_rest
 end
-@noinline function _bind_to_std_with_rest(::Type{S}, ::Bind, ::AbstractArray, ::Dims) where {S}
+@noinline function _bind_to_std_with_rest(::Type{S}, ::Bind, ::AbstractArray, ::SizeLike) where {S}
     throw(ArgumentError("Batches of variate streams containing binds must be consumed stream by stream"))
 end
diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl
index b5c54d12..3e85a3ae 100644
--- a/src/combinators/combined.jl
+++ b/src/combinators/combined.jl
@@ -215,10 +215,10 @@ end
 
 # Several variates per stream interleave the component parts, so the rows
 # of each variate are split by the fixed component sizes:
-function batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray, sz::Dims)
+function batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray, sz::SizeLike)
     n_a, n_b = _fixed_stream_length(μ.α), _fixed_stream_length(μ.β)
-    X_μ, X_rest = _batched_split(X, (n_a + n_b) * prod(sz))
-    X_v = reshape(X_μ, (n_a + n_b, sz..., Base.tail(size(X_μ))...))
+    X_μ, X_rest = _batched_split(X, _chunk_rows(n_a + n_b, sz))
+    X_v = maybestatic_reshape(X_μ, (n_a + n_b, size_dims(sz)..., Base.tail(_batch_dims(X_μ))...))
     X_a, X_b = _batched_split(X_v, n_a)
     ℓ_a, _ = batched_logdensityof_with_rest(μ.α, X_a, ())
     ℓ_b, _ = batched_logdensityof_with_rest(μ.β, X_b, ())
@@ -226,7 +226,7 @@ function batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, X::Ab
 end
 
 @inline _fixed_stream_length(μ) = _fixed_stream_length(μ, mspace_flatsize(μ))
-@inline _fixed_stream_length(μ, sz::SizeLike) = dynamic(size2length(sz))
+@inline _fixed_stream_length(μ, sz::SizeLike) = size2length(sz)
 @noinline function _fixed_stream_length(μ, ::NoMSpaceElementSize)
     throw(ArgumentError("Consuming several variates per stream requires measures of type $(nameof(typeof(μ))) to have a known variate size"))
 end
@@ -257,23 +257,23 @@ rand_impl(ctx::GenContext, μ::CombinedMeasure) = _combine_variates(μ.f_c, rand
 @inline _flat_stream(x::AbstractArray) = reduce(vcat, map(_flat_stream, x))
 @inline _flat_stream(x::Union{Tuple,NamedTuple}) = reduce(vcat, map(_flat_stream, values(x)))
 
-batched_rand_impl(ctx::GenContext, μ::CombinedMeasure, sz::Dims) = _batched_rand_pointwise(ctx, μ, sz)
+batched_rand_impl(ctx::GenContext, μ::CombinedMeasure, sz::SizeLike) = _batched_rand_pointwise(ctx, μ, sz)
 
 # Batches of merge-combined measures merge the named tuples of batches:
-function batched_rand_impl(ctx::GenContext, μ::CombinedMeasure{typeof(merge)}, sz::Dims)
+function batched_rand_impl(ctx::GenContext, μ::CombinedMeasure{typeof(merge)}, sz::SizeLike)
     merge(batched_rand_impl(ctx, μ.α, sz), batched_rand_impl(ctx, μ.β, sz))
 end
 
 # Batches of vcat-combined measures are concatenated along the streams:
-function batched_rand_impl(ctx::GenContext, μ::CombinedMeasure{typeof(vcat)}, sz::Dims)
+function batched_rand_impl(ctx::GenContext, μ::CombinedMeasure{typeof(vcat)}, sz::SizeLike)
     _combined_batched_rand(ctx, μ, sz, fixed_stream_size(μ))
 end
-function _combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::Dims, ::True)
+function _combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::SizeLike, ::True)
     A = _as_stream_batch(batched_rand_impl(ctx, μ.α, sz), μ.α)
     B = _as_stream_batch(batched_rand_impl(ctx, μ.β, sz), μ.β)
     return vcat(A, B)
 end
-_combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::Dims, ::False) = _batched_rand_pointwise(ctx, μ, sz)
+_combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::SizeLike, ::False) = _batched_rand_pointwise(ctx, μ, sz)
 
 
 # Transport consumes the variate parts of both component measures in a
@@ -339,7 +339,7 @@ function _combined_batched_to_std(::Type{S}, μ::CombinedMeasure, X::AbstractArr
     stacked(map(x -> _combined_batched_to_std(S, μ, x, static(true)), sliced(X, Val(1))))
 end
 
-function batched_transport_to_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray, sz::Dims) where {S<:StdMeasure}
+function batched_transport_to_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure}
     _combined_to_std_with_rest(S, μ, X, sz)
 end
 function _combined_to_std_with_rest(::Type{S}, μ::CombinedMeasure, X::AbstractArray, ::Tuple{}) where {S}
@@ -350,10 +350,11 @@ end
 
 # Several variates per stream interleave the component parts, so the rows
 # of each variate are split by the fixed component sizes:
-function _combined_to_std_with_rest(::Type{S}, μ::CombinedMeasure, X::AbstractArray, sz::Dims) where {S}
+function _combined_to_std_with_rest(::Type{S}, μ::CombinedMeasure, X::AbstractArray, sz::SizeLike) where {S}
     n_rows = _fixed_stream_length(μ.α) + _fixed_stream_length(μ.β)
-    X_μ, X_rest = _batched_split(X, n_rows * prod(sz))
-    Z, _ = _combined_to_std_with_rest(S, μ, reshape(X_μ, (n_rows, sz..., Base.tail(size(X_μ))...)), ())
+    X_μ, X_rest = _batched_split(X, _chunk_rows(n_rows, sz))
+    X_v = maybestatic_reshape(X_μ, (n_rows, size_dims(sz)..., Base.tail(_batch_dims(X_μ))...))
+    Z, _ = _combined_to_std_with_rest(S, μ, X_v, ())
     return _merge_multiplicity(Z, sz), X_rest
 end
 
@@ -372,7 +373,7 @@ function _combined_batched_from_std(::Type{S}, μ::CombinedMeasure, Z::AbstractA
     stacked(map(z -> transport_from_std(S, μ, z), sliced(Z, Val(1))))
 end
 
-function batched_transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, Z::AbstractArray, sz::Dims) where {S<:StdMeasure}
+function batched_transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure}
     _combined_from_std_with_rest(S, μ, Z, sz)
 end
 # Single streams yield a variate via the point protocol:
@@ -392,7 +393,7 @@ function _combined_batch_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::A
     results = map(z -> transport_from_std_with_rest(S, μ, z), sliced(Z, Val(1)))
     return stacked(map(first, results)), stacked(map(last, results))
 end
-function _combined_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, sz::Dims) where {S}
+function _combined_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, sz::SizeLike) where {S}
     _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ))
 end
 
diff --git a/src/combinators/half.jl b/src/combinators/half.jl
index 20ee71c0..9578366d 100644
--- a/src/combinators/half.jl
+++ b/src/combinators/half.jl
@@ -23,7 +23,7 @@ unhalf(μ::Half) = μ.parent
 end
 
 @inline rand_impl(ctx::GenContext, μ::Half) = abs(rand_impl(ctx, unhalf(μ)))
-@inline batched_rand_impl(ctx::GenContext, μ::Half, sz::Dims) = abs.(batched_rand_impl(ctx, unhalf(μ), sz))
+@inline batched_rand_impl(ctx::GenContext, μ::Half, sz::SizeLike) = abs.(batched_rand_impl(ctx, unhalf(μ), sz))
 
 function logdensityof_impl(μ::Half, x)
     ld = logdensityof(unhalf(μ), x) - loghalf
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 1bc9ba4b..008356be 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -66,13 +66,13 @@ function _pwr_rand(ctx::GenContext, μ::PowerMeasure, ::False)
     map(_ -> rand_impl(ctx, ν), _cartidxs(pwr_axes(μ)))
 end
 
-function batched_rand_impl(ctx::GenContext, μ::PowerMeasure, sz::Dims)
+function batched_rand_impl(ctx::GenContext, μ::PowerMeasure, sz::SizeLike)
     _pwr_batched_rand(ctx, μ, sz, fixed_stream_size(pwr_base(μ)))
 end
-function _pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::Dims, ::True)
-    batched_rand_impl(ctx, pwr_base(μ), (asnonstatic(pwr_size(μ))..., sz...))
+function _pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::SizeLike, ::True)
+    batched_rand_impl(ctx, pwr_base(μ), (size_dims(pwr_size(μ))..., size_dims(sz)...))
 end
-_pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::Dims, ::False) = _batched_rand_pointwise(ctx, μ, sz)
+_pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::SizeLike, ::False) = _batched_rand_pointwise(ctx, μ, sz)
 
 marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes)
 
@@ -160,7 +160,7 @@ _powered_stream_kernel(μ::PowerMeasure, X::AbstractArray, ::False) = _streamwis
     end
     return nothing
 end
-@inline _check_pwr_dims(::AbstractArray, ::NoMSpaceElementSize, ::Dims, ::Bool) = nothing
+@inline _check_pwr_dims(::AbstractArray, ::NoMSpaceElementSize, ::SizeLike, ::Bool) = nothing
 @inline batched_logdensityof_impl(μ::PowerMeasure, X) = _powered_kernel(logdensityof_impl, μ, X)
 @inline batched_logdensity_def(μ::PowerMeasure, X) = _powered_kernel(logdensity_def, μ, X)
 
@@ -199,14 +199,14 @@ end
 # Streams: a power consumes its size times the variates of the base measure
 # and sums the base results over its axes. Bases without fixed variate
 # sizes are consumed element by element, for single streams.
-function batched_logdensityof_with_rest(μ::PowerMeasure, X::AbstractArray, sz::Dims)
+function batched_logdensityof_with_rest(μ::PowerMeasure, X::AbstractArray, sz::SizeLike)
     _powered_ld_with_rest(μ, X, sz, fixed_stream_size(pwr_base(μ)))
 end
 function batched_logdensityof_with_rest(μ::PowerMeasure, x::AbstractVector, sz::Tuple{})
     _powered_ld_with_rest(μ, x, sz, fixed_stream_size(pwr_base(μ)))
 end
-function _powered_ld_with_rest(μ::PowerMeasure, X::AbstractArray, sz::Dims, ::True)
-    ℓ, X_rest = batched_logdensityof_with_rest(pwr_base(μ), X, (asnonstatic(pwr_size(μ))..., sz...))
+function _powered_ld_with_rest(μ::PowerMeasure, X::AbstractArray, sz::SizeLike, ::True)
+    ℓ, X_rest = batched_logdensityof_with_rest(pwr_base(μ), X, (size_dims(pwr_size(μ))..., size_dims(sz)...))
     return sum_leading_dims(ℓ, static(length(pwr_axes(μ)))), X_rest
 end
 function _powered_ld_with_rest(μ::PowerMeasure, x::AbstractVector, ::Tuple{}, ::False)
@@ -219,7 +219,7 @@ function _powered_ld_with_rest(μ::PowerMeasure, x::AbstractVector, ::Tuple{}, :
     end
     return ℓ, x_rest
 end
-@noinline function _powered_ld_with_rest(μ::PowerMeasure, ::AbstractArray, ::Dims, ::False)
+@noinline function _powered_ld_with_rest(μ::PowerMeasure, ::AbstractArray, ::SizeLike, ::False)
     throw(ArgumentError("Batches of variate streams containing powers of measures of type $(nameof(typeof(pwr_base(μ)))) must be consumed stream by stream"))
 end
 
@@ -275,7 +275,7 @@ end
     return nothing
 end
 @inline _check_pwr_flat(x::AbstractArray, k::StaticInteger, dims::Dims) = _check_pwr_dims(x, k, dims, true)
-@inline _check_pwr_flat(::AbstractArray, ::NoMSpaceElementSize, ::Dims) = _throw_size_mismatch()
+@inline _check_pwr_flat(::AbstractArray, ::NoMSpaceElementSize, ::SizeLike) = _throw_size_mismatch()
 
 checked_arg(μ::PowerMeasure, x::Any) = _throw_size_mismatch()
 
@@ -342,17 +342,17 @@ end
 # Streams: a power consumes the variates of its base measure with its size
 # as additional multiplicity. Bases without fixed variate sizes are
 # consumed element by element, for single streams.
-function batched_transport_to_std_with_rest(::Type{S}, μ::PowerMeasure, X::AbstractArray, sz::Dims) where {S<:StdMeasure}
+function batched_transport_to_std_with_rest(::Type{S}, μ::PowerMeasure, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure}
     _pwr_to_std_with_rest(S, μ, X, sz, fixed_stream_size(pwr_base(μ)))
 end
-@inline function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, X::AbstractArray, sz::Dims, ::True) where {S}
-    batched_transport_to_std_with_rest(S, pwr_base(μ), X, (asnonstatic(pwr_size(μ))..., sz...))
+@inline function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, X::AbstractArray, sz::SizeLike, ::True) where {S}
+    batched_transport_to_std_with_rest(S, pwr_base(μ), X, (size_dims(pwr_size(μ))..., size_dims(sz)...))
 end
 function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::Tuple{}, ::False) where {S}
     z, _, x_rest = transport_to_std_with_rest(S, μ, x)
     return z, x_rest
 end
-@noinline function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, ::AbstractArray, ::Dims, ::False) where {S}
+@noinline function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, ::AbstractArray, ::SizeLike, ::False) where {S}
     throw(ArgumentError("Batches of variate streams containing powers of measures of type $(nameof(typeof(pwr_base(μ)))) must be consumed stream by stream"))
 end
 
@@ -396,7 +396,7 @@ end
 
 # The stream length of a power with a base of fixed stream length:
 @inline function _fixed_stream_length(μ::PowerMeasure)
-    _fixed_stream_length(pwr_base(μ)) * prod(asnonstatic(pwr_size(μ)))
+    _fixed_stream_length(pwr_base(μ)) * size2length(pwr_size(μ))
 end
 
 # The nested variate layout of a power over its flat storage, batches of
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index ae8b05d0..b99411f3 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -78,7 +78,7 @@ _array_product_rand(ctx::GenContext, d::ProductMeasure, ::Val{false}) = _map(Bas
 
 # Batches of tuple and named tuple products are tuples resp. named tuples
 # of marginal batches:
-function batched_rand_impl(ctx::GenContext, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, sz::Dims)
+function batched_rand_impl(ctx::GenContext, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, sz::SizeLike)
     map(m -> batched_rand_impl(ctx, m, sz), marginals(μ))
 end
 
@@ -424,13 +424,13 @@ function transport_to_std_with_rest(::Type{S}, μ::ProductMeasure{<:Union{Tuple,
     return z, x_μ, x_rest
 end
 
-function batched_transport_to_std_with_rest(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, X::AbstractArray, sz::Dims) where {S<:StdMeasure}
+function batched_transport_to_std_with_rest(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure}
     _tuple_product_to_std_with_rest(S, μ, X, sz)
 end
 function _tuple_product_to_std_with_rest(::Type{S}, μ, X::AbstractArray, ::Tuple{}) where {S}
     _marginals_to_std_with_rest(S, values(marginals(μ)), X)
 end
-function _tuple_product_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::Dims) where {S}
+function _tuple_product_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::SizeLike) where {S}
     X_v, X_rest = _split_stream_variates(μ, X, sz)
     Z, _ = _marginals_to_std_with_rest(S, values(marginals(μ)), X_v)
     return _merge_multiplicity(Z, sz), X_rest
@@ -460,11 +460,11 @@ function batched_transport_from_std(::Type{S}, μ::ProductMeasure{<:Union{Tuple,
     return X
 end
 
-function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:Tuple}, Z::AbstractArray, sz::Dims) where {S<:StdMeasure}
+function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:Tuple}, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure}
     _tuple_product_from_std_with_rest(S, μ, Z, sz)
 end
 
-function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, Z::AbstractArray, sz::Dims) where {S<:StdMeasure,names}
+function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure,names}
     Xs, Z_rest = _tuple_product_from_std_with_rest(S, productmeasure(values(marginals(μ))), Z, sz)
     return NamedTuple{names}(Xs), Z_rest
 end
@@ -474,7 +474,7 @@ end
 function _tuple_product_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, ::Tuple{}) where {S}
     _marginals_batched_from_std_with_rest(S, marginals(μ), Z)
 end
-function _tuple_product_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::Dims) where {S}
+function _tuple_product_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::SizeLike) where {S}
     _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ))
 end
 
@@ -567,19 +567,19 @@ end
     throw(ArgumentError("Batched transport to products over arrays of marginals of type $(nameof(M)) requires MeasureBase.mspace_ndims to be declared for that type"))
 end
 
-function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, Z::AbstractArray, sz::Dims) where {S<:StdMeasure,M}
+function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure,M}
     _array_product_batched_from_std_with_rest(S, μ, Z, sz, _fused_marginals(M), _static_ndims_of(mspace_ndims(M)))
 end
-function _array_product_batched_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::Dims, ::Val{true}, ::Any) where {S}
+function _array_product_batched_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::SizeLike, ::Val{true}, ::Any) where {S}
     _batched_from_std_bydof(S, μ, Z, sz, length(marginals(μ)))
 end
 function _array_product_batched_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, ::Tuple{}, ::Val{false}, ::StaticInteger{K}) where {S,K}
     _marginals_from_std_loop(S, marginals(μ), Z, Val(K))
 end
-function _array_product_batched_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::Dims, ::Val{false}, ::StaticInteger{K}) where {S,K}
+function _array_product_batched_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::SizeLike, ::Val{false}, ::StaticInteger{K}) where {S,K}
     _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ))
 end
-@noinline function _array_product_batched_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, ::AbstractArray, ::Dims, ::Val{false}, ::NoMSpaceElementSize) where {S,M}
+@noinline function _array_product_batched_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, ::AbstractArray, ::SizeLike, ::Val{false}, ::NoMSpaceElementSize) where {S,M}
     throw(ArgumentError("Batched transport to products over arrays of marginals of type $(nameof(M)) requires MeasureBase.mspace_ndims to be declared for that type"))
 end
 
@@ -654,21 +654,21 @@ end
 function batched_logdensityof_with_rest(μ::ProductMeasure{<:Tuple}, x::AbstractVector, ::Tuple{})
     _marginals_ld_with_rest(marginals(μ), x)
 end
-function batched_logdensityof_with_rest(μ::ProductMeasure{<:Tuple}, X::AbstractArray, sz::Dims)
+function batched_logdensityof_with_rest(μ::ProductMeasure{<:Tuple}, X::AbstractArray, sz::SizeLike)
     X_v, X_rest = _split_stream_variates(μ, X, sz)
     ℓ, _ = _marginals_ld_with_rest(marginals(μ), X_v)
     return ℓ, X_rest
 end
-function batched_logdensityof_with_rest(μ::ProductMeasure{<:NamedTuple{names}}, X::AbstractArray, sz::Dims) where {names}
+function batched_logdensityof_with_rest(μ::ProductMeasure{<:NamedTuple{names}}, X::AbstractArray, sz::SizeLike) where {names}
     batched_logdensityof_with_rest(productmeasure(values(marginals(μ))), X, sz)
 end
 
 # The rows of `prod(sz)` variates of fixed stream length, as a batch of
 # streams `(stream length, sz..., batch dims...)`:
-function _split_stream_variates(μ, X::AbstractArray, sz::Dims)
+function _split_stream_variates(μ, X::AbstractArray, sz::SizeLike)
     n_rows = _fixed_stream_length(μ)
-    X_μ, X_rest = _batched_split(X, n_rows * prod(sz))
-    return reshape(X_μ, (n_rows, sz..., Base.tail(size(X_μ))...)), X_rest
+    X_μ, X_rest = _batched_split(X, _chunk_rows(n_rows, sz))
+    return maybestatic_reshape(X_μ, (n_rows, size_dims(sz)..., Base.tail(_batch_dims(X_μ))...)), X_rest
 end
 
 @inline _fixed_stream_length(μ::ProductMeasure{<:Tuple}) = sum(_fixed_stream_length, marginals(μ))
diff --git a/src/combinators/spikemixture.jl b/src/combinators/spikemixture.jl
index 3f05b125..02e8e57a 100644
--- a/src/combinators/spikemixture.jl
+++ b/src/combinators/spikemixture.jl
@@ -37,14 +37,14 @@ function rand_impl(ctx::GenContext, μ::SpikeMixture)
     return (rand(get_rng(ctx), get_precision(ctx)) < μ.w) * rand_impl(ctx, μ.m)
 end
 
-function batched_rand_impl(ctx::GenContext, μ::SpikeMixture, sz::Dims)
+function batched_rand_impl(ctx::GenContext, μ::SpikeMixture, sz::SizeLike)
     _spike_batched_rand(ctx, μ, sz, _static_ndims(μ.m))
 end
-function _spike_batched_rand(ctx::GenContext, μ::SpikeMixture, sz::Dims, k::StaticInteger)
+function _spike_batched_rand(ctx::GenContext, μ::SpikeMixture, sz::SizeLike, k::StaticInteger)
     X = batched_rand_impl(ctx, μ.m, sz)
     return ifelse.(_batch_mask(_rand_bulk(ctx, sz) .< μ.w, k), X, zero(eltype(X)))
 end
-_spike_batched_rand(ctx::GenContext, μ::SpikeMixture, sz::Dims, ::NoMSpaceElementSize) =
+_spike_batched_rand(ctx::GenContext, μ::SpikeMixture, sz::SizeLike, ::NoMSpaceElementSize) =
     _batched_rand_pointwise(ctx, μ, sz)
 
 testvalue(::Type{T}, μ::SpikeMixture) where {T} = zero(T)
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index 53e6fde7..aaa7d0c3 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -148,10 +148,10 @@ end
 
 # Batches of superpositions draw a batch from each component and select
 # by mass, branch-free:
-function batched_rand_impl(ctx::GenContext, μ::SuperpositionMeasure, sz::Dims)
+function batched_rand_impl(ctx::GenContext, μ::SuperpositionMeasure, sz::SizeLike)
     _superpose_batched_rand(ctx, μ, sz, _static_ndims(μ))
 end
-function _superpose_batched_rand(ctx::GenContext, μ::SuperpositionMeasure, sz::Dims, k::StaticInteger)
+function _superpose_batched_rand(ctx::GenContext, μ::SuperpositionMeasure, sz::SizeLike, k::StaticInteger)
     components = values(μ.components)
     masses, total = _component_masses(μ)
     thresholds = _batch_mask(_rand_bulk(ctx, sz) .* total, k)
@@ -163,7 +163,7 @@ function _superpose_batched_rand(ctx::GenContext, μ::SuperpositionMeasure, sz::
     end
     return X
 end
-_superpose_batched_rand(ctx::GenContext, μ::SuperpositionMeasure, sz::Dims, ::NoMSpaceElementSize) =
+_superpose_batched_rand(ctx::GenContext, μ::SuperpositionMeasure, sz::SizeLike, ::NoMSpaceElementSize) =
     _batched_rand_pointwise(ctx, μ, sz)
 
 @inline function insupport(d::SuperpositionMeasure, x)
diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl
index 1b5d7239..84aec842 100644
--- a/src/combinators/transformedmeasure.jl
+++ b/src/combinators/transformedmeasure.jl
@@ -291,13 +291,13 @@ rand_impl(ctx::GenContext, ν::PushforwardMeasure) = ν.f(rand_impl(ctx, ν.orig
 
 # Batches of pushforwards apply the function to the variates of a batch of
 # the origin:
-function batched_rand_impl(ctx::GenContext, ν::PushforwardMeasure, sz::Dims)
+function batched_rand_impl(ctx::GenContext, ν::PushforwardMeasure, sz::SizeLike)
     _pushfwd_batched_rand(ctx, ν, sz, _static_ndims(ν.origin))
 end
-@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, k::StaticInteger)
+@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::SizeLike, k::StaticInteger)
     _apply_batched(ν.f, batched_rand_impl(ctx, ν.origin, sz), k)
 end
-@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::Dims, ::NoMSpaceElementSize)
+@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::SizeLike, ::NoMSpaceElementSize)
     _batched_rand_pointwise(ctx, ν, sz)
 end
 
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
index 948adf05..f7d3b979 100644
--- a/src/combinators/weighted.jl
+++ b/src/combinators/weighted.jl
@@ -34,7 +34,7 @@ end
 end
 
 @inline rand_impl(ctx::GenContext, μ::AbstractWeightedMeasure) = rand_impl(ctx, basemeasure(μ))
-@inline batched_rand_impl(ctx::GenContext, μ::AbstractWeightedMeasure, sz::Dims) =
+@inline batched_rand_impl(ctx::GenContext, μ::AbstractWeightedMeasure, sz::SizeLike) =
     batched_rand_impl(ctx, basemeasure(μ), sz)
 
 testvalue(::Type{T}, μ::AbstractWeightedMeasure) where {T} = testvalue(T, basemeasure(μ))
@@ -89,9 +89,9 @@ insupport(μ::WeightedMeasure, x) = insupport(μ.base, x)
     batched_transport_to_std(S, basemeasure(μ), X)
 @inline batched_transport_from_std(::Type{S}, μ::AbstractWeightedMeasure, Z::AbstractArray) where {S<:StdMeasure} =
     batched_transport_from_std(S, basemeasure(μ), Z)
-@inline batched_transport_to_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, X::AbstractArray, sz::Dims) where {S<:StdMeasure} =
+@inline batched_transport_to_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure} =
     batched_transport_to_std_with_rest(S, basemeasure(μ), X, sz)
-@inline batched_transport_from_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, Z::AbstractArray, sz::Dims) where {S<:StdMeasure} =
+@inline batched_transport_from_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure} =
     batched_transport_from_std_with_rest(S, basemeasure(μ), Z, sz)
 
 Adapt.adapt_structure(to, μ::WeightedMeasure) = WeightedMeasure(μ.logweight, Adapt.adapt(to, μ.base))
diff --git a/src/density-batched.jl b/src/density-batched.jl
index 61e9f6a0..e50d555d 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -229,7 +229,7 @@ end
 # streams, batched as `(rows, batch dims...)`.
 
 """
-    MeasureBase.batched_logdensityof_with_rest(μ::AbstractMeasure, X, sz::Dims)
+    MeasureBase.batched_logdensityof_with_rest(μ::AbstractMeasure, X, sz::SizeLike)
 
 Consume variates of `μ` from the batch `X` of flat vector streams (first
 dimension along the streams, further dimensions are batch dimensions), a
@@ -250,7 +250,7 @@ combinators consume batches stream by stream.
 """
 function batched_logdensityof_with_rest end
 
-function batched_logdensityof_with_rest(μ::AbstractMeasure, X::AbstractArray, sz::Dims)
+function batched_logdensityof_with_rest(μ::AbstractMeasure, X::AbstractArray, sz::SizeLike)
     _stream_ld_with_rest(logdensityof_impl, μ, X, sz)
 end
 
@@ -260,7 +260,7 @@ function batched_logdensityof_with_rest(μ::AbstractMeasure, x::AbstractVector,
     return ℓ, x_rest
 end
 
-function _stream_ld_with_rest(f::F, μ, X::AbstractArray, sz::Dims) where {F}
+function _stream_ld_with_rest(f::F, μ, X::AbstractArray, sz::SizeLike) where {F}
     vsz = _stream_consume_size(μ)
     X_μ, X_rest = _batched_consume(X, vsz, sz)
     return _consumed_ld(f, μ, X_μ, vsz), X_rest
@@ -275,15 +275,14 @@ end
 # batch of streams as a flat batch `(vsz..., sz..., batch dims...)`; scalar
 # variates as `(1, sz..., batch dims...)`. Static sizes keep static
 # streams static.
-@inline function _batched_consume(X::AbstractArray, vsz::SizeLike, sz::Dims)
+@inline function _batched_consume(X::AbstractArray, vsz::SizeLike, sz::SizeLike)
     dims = _consumed_dims(vsz)
     X_flat, X_rest = _batched_split(X, _chunk_rows(prod(dims), sz))
     return _reshape_consumed(X_flat, (dims..., sz...)), X_rest
 end
 @inline _consumed_dims(::Tuple{}) = (static(1),)
 @inline _consumed_dims(vsz::SizeLike) = size_dims(vsz)
-@inline _chunk_rows(n::IntegerLike, ::Tuple{}) = n
-@inline _chunk_rows(n::IntegerLike, sz::Dims) = dynamic(n) * prod(sz)
+@inline _chunk_rows(n::IntegerLike, sz::SizeLike) = n * size2length(sz)
 
 @inline _reshape_consumed(X_flat::AbstractArray, ::Tuple{IntegerLike}) = X_flat
 @inline function _reshape_consumed(X_flat::AbstractArray, dims::Tuple{Vararg{IntegerLike}})
diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl
index 4fdbb610..dea1c087 100644
--- a/src/primitives/dirac.jl
+++ b/src/primitives/dirac.jl
@@ -35,7 +35,7 @@ logdensity_def(::Dirac, x::Number) = zero(float(typeof(x)))
 logdensity_def(::Dirac, x) = zero(_logd_numtype(x))
 
 @inline rand_impl(::GenContext, μ::Dirac) = μ.x
-@inline batched_rand_impl(ctx::GenContext, μ::Dirac, sz::Dims) = _const_batch(ctx, μ.x, sz)
+@inline batched_rand_impl(ctx::GenContext, μ::Dirac, sz::SizeLike) = _const_batch(ctx, μ.x, sz)
 
 export dirac
 
diff --git a/src/rand.jl b/src/rand.jl
index a1b48063..ec0e4499 100644
--- a/src/rand.jl
+++ b/src/rand.jl
@@ -30,12 +30,13 @@ Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractMeasure) where {T<:AbstractFl
 
 
 """
-    MeasureBase.batched_rand_impl(ctx::GenContext, μ, sz::Dims)
+    MeasureBase.batched_rand_impl(ctx::GenContext, μ, sz::SizeLike)
 
 Generate a batch of random variates of `μ` of batch size `sz` in flat
 form, an array `(variate dims..., sz...)`, or a single variate for
 `sz == ()`. Batches of tuple and named tuple variates are tuples resp.
-named tuples of batches.
+named tuples of batches. Fully static batch sizes give static arrays on
+the CPU.
 
 This is the primary extension point for random variate generation. The
 default implementation draws a batch of variates of the preferred
@@ -63,30 +64,30 @@ struct _NoRandImpl end
 struct _MaybeRandImpl end
 
 @inline rand_impl(ctx::GenContext, μ) = _rand_default(ctx, μ, (), _NoRandImpl())
-@inline batched_rand_impl(ctx::GenContext, μ, sz::Dims) = _rand_default(ctx, μ, sz, _MaybeRandImpl())
+@inline batched_rand_impl(ctx::GenContext, μ, sz::SizeLike) = _rand_default(ctx, μ, sz, _MaybeRandImpl())
 
-@inline _rand_default(ctx::GenContext, μ, sz::Dims, m) = _rand_via_std(ctx, μ, sz, preferred_stdmeasure(μ), m)
+@inline _rand_default(ctx::GenContext, μ, sz::SizeLike, m) = _rand_via_std(ctx, μ, sz, preferred_stdmeasure(μ), m)
 
-@inline function _rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{S}, m) where {S<:StdMeasure}
+@inline function _rand_via_std(ctx::GenContext, μ, sz::SizeLike, ::Type{S}, m) where {S<:StdMeasure}
     _rand_via_std_dof(ctx, μ, sz, S, fast_dof(μ), m)
 end
-@inline _rand_via_std(ctx::GenContext, μ, sz::Dims, ::Type{AnyStdMeasure}, m) = _rand_via_std(ctx, μ, sz, StdUniform, m)
-@inline _rand_via_std(ctx::GenContext, μ, sz::Dims, ::Any, m) = _rand_pointwise(ctx, μ, sz, m)
+@inline _rand_via_std(ctx::GenContext, μ, sz::SizeLike, ::Type{AnyStdMeasure}, m) = _rand_via_std(ctx, μ, sz, StdUniform, m)
+@inline _rand_via_std(ctx::GenContext, μ, sz::SizeLike, ::Any, m) = _rand_pointwise(ctx, μ, sz, m)
 
-function _rand_via_std_dof(ctx::GenContext, μ, sz::Dims, ::Type{S}, n::IntegerLike, ::Any) where {S<:StdMeasure}
-    convert_realtype(get_precision(ctx), batched_transport_from_std(S, μ, _rand_std(ctx, S, (dynamic(n), sz...))))
+function _rand_via_std_dof(ctx::GenContext, μ, sz::SizeLike, ::Type{S}, n::IntegerLike, ::Any) where {S<:StdMeasure}
+    convert_realtype(get_precision(ctx), batched_transport_from_std(S, μ, _rand_std(ctx, S, (n, size_dims(sz)...))))
 end
-@inline _rand_via_std_dof(ctx::GenContext, μ, sz::Dims, ::Type, ::Any, m) = _rand_pointwise(ctx, μ, sz, m)
+@inline _rand_via_std_dof(ctx::GenContext, μ, sz::SizeLike, ::Type, ::Any, m) = _rand_pointwise(ctx, μ, sz, m)
 
 # Variates generated one by one, stacked into a flat batch:
-@inline _rand_pointwise(ctx::GenContext, μ, sz::Dims, ::Any) = _batched_rand_pointwise(ctx, μ, sz)
+@inline _rand_pointwise(ctx::GenContext, μ, sz::SizeLike, ::Any) = _batched_rand_pointwise(ctx, μ, sz)
 @inline _rand_pointwise(ctx::GenContext, μ, ::Tuple{}, ::_MaybeRandImpl) = rand_impl(ctx, μ)
 @noinline function _rand_pointwise(::GenContext, μ, ::Tuple{}, ::_NoRandImpl)
     throw(ArgumentError("Random variate generation is not implemented for measures of type $(nameof(typeof(μ))), define MeasureBase.batched_rand_impl or MeasureBase.rand_impl"))
 end
 
-function _batched_rand_pointwise(ctx::GenContext, μ, sz::Dims)
-    _stack_variates(map(_ -> rand_impl(ctx, μ), CartesianIndices(sz)))
+function _batched_rand_pointwise(ctx::GenContext, μ, sz::SizeLike)
+    _stack_variates(map(_ -> rand_impl(ctx, μ), CartesianIndices(asnonstatic(sz))))
 end
 @inline _batched_rand_pointwise(ctx::GenContext, μ, ::Tuple{}) = rand_impl(ctx, μ)
 
@@ -98,14 +99,25 @@ end
 # Bulk draws of standard variates on the compute unit, single draws for
 # zero batch dimensions:
 
-@inline _rand_std(ctx::GenContext, ::Type{S}, dims::Dims) where {S<:StdMeasure} = batched_rand_impl(ctx, S(), dims)
+@inline _rand_std(ctx::GenContext, ::Type{S}, dims::SizeLike) where {S<:StdMeasure} = batched_rand_impl(ctx, S(), dims)
 
-@inline _rand_bulk(ctx::GenContext, sz::Dims) = rand(ctx, sz)
-@inline _randn_bulk(ctx::GenContext, sz::Dims) = randn(ctx, sz)
-@inline _randexp_bulk(ctx::GenContext, sz::Dims) = _randexp_bulk(ctx, sz, get_compute_unit(ctx))
-@inline _randexp_bulk(ctx::GenContext, sz::Dims, ::CPUnit) = randexp(ctx, sz)
+@inline _rand_bulk(ctx::GenContext, sz::SizeLike) = _bulk_draw(rand, ctx, sz)
+@inline _randn_bulk(ctx::GenContext, sz::SizeLike) = _bulk_draw(randn, ctx, sz)
+@inline _randexp_bulk(ctx::GenContext, sz::SizeLike) = _randexp_bulk(ctx, sz, get_compute_unit(ctx))
+@inline _randexp_bulk(ctx::GenContext, sz::SizeLike, ::CPUnit) = _bulk_draw(randexp, ctx, sz)
 # Not all compute units provide exponential draws, derive them from uniform draws then:
-@inline _randexp_bulk(ctx::GenContext, sz::Dims, ::AbstractComputeUnit) = -log1p.(-_rand_bulk(ctx, sz))
+@inline _randexp_bulk(ctx::GenContext, sz::SizeLike, ::AbstractComputeUnit) = -log1p.(-_rand_bulk(ctx, sz))
+
+# Fully static batch sizes draw static arrays on the CPU, so that variates
+# of statically sized measures are allocation-free. Other compute units
+# allocate their own arrays.
+@inline _bulk_draw(f::F, ctx::GenContext, sz::SizeLike) where {F} = f(ctx, asnonstatic(sz))
+@inline _bulk_draw(f::F, ctx::GenContext, sz::StaticSizeLike) where {F} =
+    _bulk_draw(f, ctx, sz, get_compute_unit(ctx))
+@inline _bulk_draw(f::F, ctx::GenContext, sz::StaticSizeLike, ::AbstractComputeUnit) where {F} =
+    f(ctx, asnonstatic(sz))
+@inline _bulk_draw(f::F, ctx::GenContext, sz::StaticSizeLike, ::CPUnit) where {F} =
+    f(get_rng(ctx), staticarray_type(get_precision(ctx), canonical_size(sz)))
 
 @inline _rand_bulk(ctx::GenContext, ::Tuple{}) = rand(get_rng(ctx), get_precision(ctx))
 @inline _randn_bulk(ctx::GenContext, ::Tuple{}) = randn(get_rng(ctx), get_precision(ctx))
@@ -113,13 +125,16 @@ end
 
 # Test values use a constant RNG, which only draws single values:
 const _ConstantContext = GenContext{<:AbstractFloat,<:AbstractComputeUnit,ConstantRNG}
-@inline _rand_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, rand(ConstantRNG(), get_precision(ctx)), sz)
-@inline _randn_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, randn(ConstantRNG(), get_precision(ctx)), sz)
-@inline _randexp_bulk(ctx::_ConstantContext, sz::Dims) = _const_bulk(ctx, randexp(ConstantRNG(), get_precision(ctx)), sz)
+@inline _rand_bulk(ctx::_ConstantContext, sz::SizeLike) = _const_bulk(ctx, rand(ConstantRNG(), get_precision(ctx)), sz)
+@inline _randn_bulk(ctx::_ConstantContext, sz::SizeLike) = _const_bulk(ctx, randn(ConstantRNG(), get_precision(ctx)), sz)
+@inline _randexp_bulk(ctx::_ConstantContext, sz::SizeLike) = _const_bulk(ctx, randexp(ConstantRNG(), get_precision(ctx)), sz)
 @inline _rand_bulk(ctx::_ConstantContext, ::Tuple{}) = rand(ConstantRNG(), get_precision(ctx))
 @inline _randn_bulk(ctx::_ConstantContext, ::Tuple{}) = randn(ConstantRNG(), get_precision(ctx))
 @inline _randexp_bulk(ctx::_ConstantContext, ::Tuple{}) = randexp(ConstantRNG(), get_precision(ctx))
-@inline _const_bulk(ctx::GenContext, x, sz::Dims) = fill!(allocate_array(ctx, typeof(x), sz), x)
+@inline _const_bulk(ctx::GenContext, x, sz::SizeLike) = _const_bulk(ctx, x, sz, get_compute_unit(ctx))
+@inline _const_bulk(ctx::GenContext, x, sz::SizeLike, ::AbstractComputeUnit) =
+    fill!(allocate_array(ctx, typeof(x), asnonstatic(sz)), x)
+@inline _const_bulk(ctx::GenContext, x, sz::StaticSizeLike, ::CPUnit) = maybestatic_fill(x, sz)
 
 # A mask over the batch dimensions, aligned with a flat batch of variates
 # of rank `k`:
@@ -130,10 +145,10 @@ const _ConstantContext = GenContext{<:AbstractFloat,<:AbstractComputeUnit,Consta
 end
 
 # A batch of copies of a constant variate:
-function _const_batch(ctx::GenContext, x, sz::Dims)
-    X = allocate_array(ctx, eltype(x), (size(x)..., sz...))
+function _const_batch(ctx::GenContext, x, sz::SizeLike)
+    X = allocate_array(ctx, eltype(x), (size(x)..., asnonstatic(sz)...))
     X .= x
     return X
 end
-@inline _const_batch(ctx::GenContext, x::Number, sz::Dims) = fill!(allocate_array(ctx, typeof(x), sz), x)
+@inline _const_batch(ctx::GenContext, x::Number, sz::SizeLike) = _const_bulk(ctx, x, sz)
 @inline _const_batch(::GenContext, x::Number, ::Tuple{}) = x
diff --git a/src/standard/stdexponential.jl b/src/standard/stdexponential.jl
index a9823084..124ce62a 100644
--- a/src/standard/stdexponential.jl
+++ b/src/standard/stdexponential.jl
@@ -22,4 +22,4 @@ end
 @inline transport_def(::StdExponential, μ::StdUniform, x) = _nan_outside(μ, x, -log1p(-_unit_interior(x)))
 
 @inline rand_impl(ctx::GenContext, ::StdExponential) = randexp(get_rng(ctx), get_precision(ctx))
-@inline batched_rand_impl(ctx::GenContext, ::StdExponential, sz::Dims) = _randexp_bulk(ctx, sz)
+@inline batched_rand_impl(ctx::GenContext, ::StdExponential, sz::SizeLike) = _randexp_bulk(ctx, sz)
diff --git a/src/standard/stdlogistic.jl b/src/standard/stdlogistic.jl
index 98b76219..018beed3 100644
--- a/src/standard/stdlogistic.jl
+++ b/src/standard/stdlogistic.jl
@@ -19,7 +19,7 @@ export StdLogistic
 @inline transport_def(::StdLogistic, μ::StdUniform, p) = _nan_outside(μ, p, logit(_unit_interior(p)))
 
 @inline rand_impl(ctx::GenContext, ::StdLogistic) = logit(rand(get_rng(ctx), get_precision(ctx)))
-@inline batched_rand_impl(ctx::GenContext, ::StdLogistic, sz::Dims) = logit.(_rand_bulk(ctx, sz))
+@inline batched_rand_impl(ctx::GenContext, ::StdLogistic, sz::SizeLike) = logit.(_rand_bulk(ctx, sz))
 
 smf(::StdLogistic, x) = logistic(x)
 smf(::StdLogistic) = logistic
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index 3c2ab19b..9ee949f6 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -10,6 +10,11 @@ StdMeasure(::typeof(randn)) = StdNormal()
 
 @inline check_dof(::StdMeasure, ::StdMeasure) = nothing
 
+# Standard measures have real scalar variates, checking them directly keeps
+# the recursion over base measures (and its boxed arguments) out of the
+# kernels:
+@inline checked_arg(::StdMeasure, x::Real) = x
+
 @inline massof(::StdMeasure) = static(1.0)
 
 @inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x
diff --git a/src/standard/stdnormal.jl b/src/standard/stdnormal.jl
index 58b0408e..a94d92e4 100644
--- a/src/standard/stdnormal.jl
+++ b/src/standard/stdnormal.jl
@@ -21,7 +21,7 @@ export StdNormal
 @inline getdof(::StdNormal) = static(1)
 
 @inline rand_impl(ctx::GenContext, ::StdNormal) = randn(get_rng(ctx), get_precision(ctx))
-@inline batched_rand_impl(ctx::GenContext, ::StdNormal, sz::Dims) = _randn_bulk(ctx, sz)
+@inline batched_rand_impl(ctx::GenContext, ::StdNormal, sz::SizeLike) = _randn_bulk(ctx, sz)
 
 Φ(z) = erfc(-z * invsqrt2) / 2
 Φinv(p) = -erfcinv(2 * p) * sqrt2
diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl
index 3c45f078..5971cd0f 100644
--- a/src/standard/stduniform.jl
+++ b/src/standard/stduniform.jl
@@ -19,7 +19,7 @@ end
 @inline basemeasure(::StdUniform) = LebesgueBase()
 
 @inline rand_impl(ctx::GenContext, ::StdUniform) = rand(get_rng(ctx), get_precision(ctx))
-@inline batched_rand_impl(ctx::GenContext, ::StdUniform, sz::Dims) = _rand_bulk(ctx, sz)
+@inline batched_rand_impl(ctx::GenContext, ::StdUniform, sz::SizeLike) = _rand_bulk(ctx, sz)
 
 massof(::StdUniform, s::Interval) = massof(Lebesgue(0.0 .. 1.0), s)
 
diff --git a/src/transport-batched.jl b/src/transport-batched.jl
index 79a1e993..5878e865 100644
--- a/src/transport-batched.jl
+++ b/src/transport-batched.jl
@@ -118,7 +118,7 @@ end
 
 
 """
-    MeasureBase.batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::Dims)
+    MeasureBase.batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::SizeLike)
 
 Consume variates of `μ` from the batch `X` of flat vector streams (first
 dimension along the streams, further dimensions are batch dimensions), a
@@ -136,7 +136,7 @@ whose variates are composed of the variates of other measures implement
 """
 function batched_transport_to_std_with_rest end
 
-function batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::Dims) where {S<:StdMeasure}
+function batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure}
     _to_std_with_rest_default(S, μ, X, sz)
 end
 
@@ -144,7 +144,7 @@ function _to_std_with_rest_default(::Type{S}, μ, x::AbstractVector, ::Tuple{})
     z, _, x_rest = transport_to_std_with_rest(S, μ, x)
     return z, x_rest
 end
-function _to_std_with_rest_default(::Type{S}, μ, X::AbstractArray, sz::Dims) where {S}
+function _to_std_with_rest_default(::Type{S}, μ, X::AbstractArray, sz::SizeLike) where {S}
     vsz = _stream_consume_size(μ)
     X_μ, X_rest = _batched_consume(X, vsz, sz)
     Z = batched_transport_to_std(S, μ, _consumed_variates(X_μ, vsz))
@@ -156,15 +156,15 @@ end
 # Standard variates of `prod(sz)` variates per stream, `(dof, sz..., batch
 # dims...)`, as one stream chunk `(dof * prod(sz), batch dims...)`, and
 # back:
-@inline _merge_multiplicity(Z::AbstractArray, sz::Dims) = merge_leading_dims(Z, static(1) + static(length(sz)))
+@inline _merge_multiplicity(Z::AbstractArray, sz::SizeLike) = merge_leading_dims(Z, static(1) + maybestatic_length(size_dims(sz)))
 @inline _split_multiplicity(Z::AbstractArray, ::Tuple{}, n) = Z
-@inline function _split_multiplicity(Z::AbstractArray, sz::Dims, n)
+@inline function _split_multiplicity(Z::AbstractArray, sz::SizeLike, n)
     maybestatic_reshape(Z, (n, sz..., Base.tail(_batch_dims(Z))...))
 end
 
 
 """
-    MeasureBase.batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::Dims)
+    MeasureBase.batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::SizeLike)
 
 Consume standard variates of type `S` for a batch of variates of size
 `sz` per stream from the batch `Z` of streams of standard variates (first
@@ -180,20 +180,20 @@ variates are composed of the variates of other measures implement
 """
 function batched_transport_from_std_with_rest end
 
-function batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::Dims) where {S<:StdMeasure}
+function batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure}
     _from_std_with_rest_default(S, μ, Z, sz)
 end
 
 _from_std_with_rest_default(::Type{S}, μ, z::AbstractVector, ::Tuple{}) where {S} = transport_from_std_with_rest(S, μ, z)
-function _from_std_with_rest_default(::Type{S}, μ, Z::AbstractArray, sz::Dims) where {S}
+function _from_std_with_rest_default(::Type{S}, μ, Z::AbstractArray, sz::SizeLike) where {S}
     _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ))
 end
 
-function _batched_from_std_bydof(::Type{S}, μ, Z::AbstractArray, sz::Dims, n::IntegerLike) where {S}
+function _batched_from_std_bydof(::Type{S}, μ, Z::AbstractArray, sz::SizeLike, n::IntegerLike) where {S}
     Z_μ, Z_rest = _batched_split(Z, _chunk_rows(n, sz))
     return batched_transport_from_std(S, μ, _split_multiplicity(Z_μ, sz, n)), Z_rest
 end
-@noinline function _batched_from_std_bydof(::Type{S}, μ, ::AbstractArray, ::Dims, ::AbstractNoDOF) where {S}
+@noinline function _batched_from_std_bydof(::Type{S}, μ, ::AbstractArray, ::SizeLike, ::AbstractNoDOF) where {S}
     throw(ArgumentError("Batched transport from standard measures requires measures of type $(nameof(typeof(μ))) to have fast degrees of freedom or to implement MeasureBase.batched_transport_from_std_with_rest"))
 end
 

From f359023d43a07a513738b368e746850e4ee53669 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 21 Sep 2026 16:25:55 +0200
Subject: [PATCH 116/122] Test static variates end to end

A hierarchical model over a named tuple of scalar and static-array
marginals (a bind of a named tuple product with a kernel whose marginals
depend on the primary variate but whose sizes don't) checks that random
variates, transports to and from a static standard power and densities
are type stable and allocation free, plus static powers, nested static
powers and batches of static variates.

`test/transport.jl` imports `MeasureBase` itself, so that it runs on its
own and not only as part of the suite.

Created by generative AI.
---
 test/runtests.jl        |   1 +
 test/static_variates.jl | 129 ++++++++++++++++++++++++++++++++++++++++
 test/transport.jl       |   2 +
 3 files changed, 132 insertions(+)
 create mode 100644 test/static_variates.jl

diff --git a/test/runtests.jl b/test/runtests.jl
index 4fdbc12d..e53acc85 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -19,6 +19,7 @@ include("shape_contract.jl")
 include("logdensities.jl")
 include("structured_batches.jl")
 include("batched_regressions.jl")
+include("static_variates.jl")
 include("fixed_size_arrays.jl")
 include("support_conventions.jl")
 include("numtype.jl")
diff --git a/test/static_variates.jl b/test/static_variates.jl
new file mode 100644
index 00000000..eeb2e3f5
--- /dev/null
+++ b/test/static_variates.jl
@@ -0,0 +1,129 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+# Static variates end to end: measures whose variate sizes are statically
+# known generate, transport and evaluate static arrays, type stable and
+# allocation free.
+
+using Test
+
+using MeasureBase
+using MeasureBase: StdNormal, StdUniform, StdExponential
+using MeasureBase: productmeasure, mbind, weightedmeasure, transport_to, logdensityof
+using MeasureBase.InverseFunctions: inverse
+using ArraysOfArrays: ArrayOfSimilarArrays, flatview, sliced
+using StaticArrays: SVector, SMatrix
+using Static: static
+
+include("testutils.jl")
+
+# A hierarchical model over a named tuple of scalar and static-array
+# marginals: the secondary marginals depend on the primary variate, their
+# sizes don't.
+const static_primary = productmeasure((
+    a = StdNormal(),
+    b = weightedmeasure(-0.5, StdExponential()),
+))
+
+static_kernel(x) = productmeasure((
+    c = StdUniform()^static(2),
+    d = weightedmeasure(-abs(x.a), StdNormal()^static(3)),
+))
+
+const static_model = mbind(static_kernel, static_primary, merge)
+
+@testset "static variates" begin
+    @testset "static powers" begin
+        @test @inferred(rand(StdNormal()^static(3))) isa SVector{3,Float64}
+        @test @inferred(rand(StdUniform()^static(2))) isa SVector{2,Float64}
+        @test @inferred(rand(StdExponential()^static(4))) isa SVector{4,Float64}
+
+        # Nested powers keep the flat `(base dims..., power dims...)` rule:
+        x = @inferred(rand((StdNormal()^static(2))^static(3)))
+        @test x isa ArrayOfSimilarArrays{Float64,1,1,<:SMatrix{2,3}}
+        @test flatview(x) isa SMatrix{2,3,Float64}
+
+        μ = StdNormal()^static(3)
+        z = SVector(0.1, 0.2, 0.3)
+        @test @inferred(transport_to(StdUniform()^static(3), μ)(z)) isa SVector{3,Float64}
+        @test @inferred(transport_to(μ, StdUniform()^static(3))(SVector(0.1, 0.5, 0.9))) isa
+              SVector{3,Float64}
+        @test @inferred(logdensityof(μ, z)) ≈ sum(logdensityof.(Ref(StdNormal()), z))
+        @test allocations_of(logdensityof, μ, z) == 0
+        @test allocations_of(transport_to(StdUniform()^static(3), μ), z) == 0
+    end
+
+    @testset "hierarchical model with static marginals" begin
+        μ = static_model
+        x = @inferred rand(μ)
+        @test x isa NamedTuple{(:a, :b, :c, :d)}
+        @test x.a isa Float64
+        @test x.b isa Float64
+        @test x.c isa SVector{2,Float64}
+        @test x.d isa SVector{3,Float64}
+
+        ν = StdNormal()^static(7)
+        f = transport_to(ν, μ)
+        f_inv = inverse(f)
+
+        z = @inferred f(x)
+        @test z isa SVector{7,Float64}
+        y = @inferred f_inv(z)
+        @test y isa NamedTuple{(:a, :b, :c, :d)}
+        @test y.c isa SVector{2,Float64}
+        @test y.d isa SVector{3,Float64}
+        @test all(map((u, v) -> u ≈ v, values(y), values(x)))
+
+        @test allocations_of(f, x) == 0
+        @test allocations_of(f_inv, z) == 0
+        @test allocations_of(logdensityof, μ, x) == 0
+
+        # The model density is the sum of the component densities:
+        x_a = (a = x.a, b = x.b)
+        x_b = (c = x.c, d = x.d)
+        @test @inferred(logdensityof(μ, x)) ≈
+              logdensityof(static_primary, x_a) + logdensityof(static_kernel(x_a), x_b)
+    end
+
+    @testset "products of scalar and static-array marginals" begin
+        μ = productmeasure((
+            a = StdNormal(),
+            b = StdUniform()^static(2),
+            c = weightedmeasure(-0.25, StdExponential()^static(3)),
+        ))
+        x = @inferred rand(μ)
+        @test x isa NamedTuple{(:a, :b, :c)}
+        @test x.a isa Float64
+        @test x.b isa SVector{2,Float64}
+        @test x.c isa SVector{3,Float64}
+
+        f = transport_to(StdNormal()^static(6), μ)
+        f_inv = inverse(f)
+        z = @inferred f(x)
+        @test z isa SVector{6,Float64}
+        y = @inferred f_inv(z)
+        @test y isa NamedTuple{(:a, :b, :c)}
+        @test all(map((u, v) -> u ≈ v, values(y), values(x)))
+        @test allocations_of(f, x) == 0
+        @test allocations_of(f_inv, z) == 0
+        @test allocations_of(logdensityof, μ, x) == 0
+
+        # Tuple products behave the same way:
+        μ_t = productmeasure((StdNormal(), StdUniform()^static(2)))
+        x_t = @inferred rand(μ_t)
+        @test x_t isa Tuple{Float64,SVector{2,Float64}}
+        f_t = transport_to(StdNormal()^static(3), μ_t)
+        @test @inferred(f_t(x_t)) isa SVector{3,Float64}
+        @test @inferred(inverse(f_t)(f_t(x_t))) isa Tuple{Float64,SVector{2,Float64}}
+    end
+
+    @testset "batches of static variates" begin
+        μ = StdNormal()^static(3)
+        X = SMatrix{3,4}(reshape(collect(1:12) ./ 10, 3, 4))
+        ℓ = logdensities(μ, X)
+        @test ℓ ≈ [logdensityof(μ, SVector{3}(X[:, i])) for i in 1:4]
+
+        ν = StdUniform()^static(3)
+        Y = transport_to(ν, μ).(sliced(X, Val(1)))
+        @test flatview(Y) ≈ reduce(hcat, [transport_to(ν, μ)(SVector{3}(X[:, i])) for i in 1:4])
+    end
+end
diff --git a/test/transport.jl b/test/transport.jl
index d2c31ae1..24734fa3 100644
--- a/test/transport.jl
+++ b/test/transport.jl
@@ -1,6 +1,8 @@
 using Test
 
+import MeasureBase
 using MeasureBase.Interface: transport_to, test_transport
+using MeasureBase: AbstractMeasure
 using MeasureBase: StdUniform, StdExponential, StdLogistic, StdNormal
 using MeasureBase: Dirac, Half, restrict, mbind, productmeasure, pushfwd
 using MeasureBase: transport_to_std, transport_from_std, transport_from_std_with_rest

From 85ce9dd5368e0f450c404280bd3300b290a40681 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 21 Sep 2026 16:26:19 +0200
Subject: [PATCH 117/122] Record static variates in the redesign notes

Created by generative AI.
---
 redesign.md | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/redesign.md b/redesign.md
index 90302dd5..1ec36c3d 100644
--- a/redesign.md
+++ b/redesign.md
@@ -77,6 +77,16 @@ entries per variate, composed measures implement the with-rest forms.
 `getdof`/`fast_dof` are declaration-derived, used at construction time
 and for chunking, never inside kernels.
 
+**Static variates.** Sizes are `StaticThings.SizeLike` throughout, the
+with-rest multiplicity and the batch size of `rand` included, so a
+statically sized power stays static through streams. On the CPU, bulk
+draws of a fully static size are static arrays and the stream splits and
+reshapes keep them static: `rand`, `transport_to` and `logdensityof` of a
+measure whose variate sizes are all static are type stable and allocation
+free, for powers, tuple and named tuple products and binds over them.
+Sizes that depend on variate values stay dynamic, and other compute units
+allocate their own arrays.
+
 **Random variates.** `rand(ctx::GenContext, μ)` with RNG, precision and
 compute unit (`rand(μ)`, `rand(rng, μ)`, `rand(T, μ)` are wrappers).
 `batched_rand_impl(ctx, μ, sz::Dims)` returns a flat batch, a single
@@ -166,6 +176,8 @@ Full suite (Aqua, extensions, doctests) on CPU with JLArrays cases and
 a FixedSizeArrays transparency check (fixed-size inputs give fixed-size
 outputs; FixedSizeArrays stays a test dependency, allocating fixed-size
 variates by default is a HeterogeneousComputing decision for later).
+`test/static_variates.jl` checks that statically sized measures stay type
+stable and allocation free.
 `test/test_reactant.jl` runs as part of the suite on 64-bit Linux and
 macOS with stable Julia, adding Reactant on demand as MGVI does (backend
 via `MEASUREBASE_REACTANT_BACKEND`). `test/cuda` is opt-in. Both run
@@ -211,8 +223,8 @@ locally on the GB10, green at HEAD except one expected-broken CUDA case
   Quantile results within `4 eps` of the support edges snap to the edges
   only on the generic logistic path of wrapped distributions.
 - Decisions pending: `Half` tails via log-ccdf, device random variate
-  infrastructure and `rand!`, Tier-1 static variates, the
-  `smart-constructors.jl` review (location-scale arrays as affine
-  pushforwards of powers), `_static_ndims` type-first vs. instance-first.
+  infrastructure and `rand!`, the `smart-constructors.jl` review
+  (location-scale arrays as affine pushforwards of powers),
+  `_static_ndims` type-first vs. instance-first.
 - Polish before merge: docs pass, NEWS, history curation, version bump,
   remove this file.

From f770668232111a1b844a82be379dc0c04ee1820b Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 21 Sep 2026 17:07:30 +0200
Subject: [PATCH 118/122] Keep variate ranks plain integers
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`mspace_ndims` returned a `StaticInt` where it followed from a variate
size and an `Int` where a measure type declared it, which made
`mspace_ndims(typeof(μ))` and `mspace_ndims(μ)` disagree and forced
`IntegerLike` on the dispatch that consumes them. Ranks are `Int`s again,
constant-folded through the type-level definitions, and staticness comes
from `_static_ndims` as before.

Follows StaticThings' rename of the mapped tuple-type fold.

Created by generative AI.
---
 src/MeasureBase.jl           |  2 +-
 src/combinators/power.jl     |  6 +++---
 src/combinators/product.jl   |  4 ++--
 src/combinators/superpose.jl |  6 +++---
 src/density-batched.jl       | 10 +++++-----
 src/mspace.jl                |  8 ++++----
 6 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
index f584b9f1..f4a72f7e 100644
--- a/src/MeasureBase.jl
+++ b/src/MeasureBase.jl
@@ -51,7 +51,7 @@ using StaticThings:
     maybestatic_first, maybestatic_last, maybestatic_view,
     maybestatic_oneto, maybestatic_fill, maybestatic_reshape,
     axes2size, size2length, split_at, staticarray_type,
-    static_all, static_any, static_reduce,
+    static_all, static_any, static_mapreduce,
     sum_leading_dims, drop_leading_dims, merge_leading_dims, all_leading_dims
 
 import HeterogeneousComputing
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 008356be..37bc3e7b 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -107,9 +107,9 @@ end
 @inline function mspace_ndims(::Type{<:PowerMeasure{M,A}}) where {M,A<:Tuple}
     _pwr_ndims(mspace_ndims(M), fieldcount(A), fixed_stream_size(M))
 end
-@inline _pwr_ndims(n::IntegerLike, k::IntegerLike, ::Any) = n + k
-@inline _pwr_ndims(::NoMSpaceElementSize, ::IntegerLike, ::True) = static(1)
-@inline _pwr_ndims(n::NoMSpaceElementSize, ::IntegerLike, ::False) = n
+@inline _pwr_ndims(n::Integer, k::Integer, ::Any) = n + k
+@inline _pwr_ndims(::NoMSpaceElementSize, ::Integer, ::True) = 1
+@inline _pwr_ndims(n::NoMSpaceElementSize, ::Integer, ::False) = n
 
 # Local measures of powers at nested variates are products of the local
 # measures of the elements:
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index b99411f3..7e202e8a 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -275,7 +275,7 @@ end
 @inline function logdensity_def(μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray{<:Number}) where {M}
     _array_product_ld(logdensity_def, μ, x, mspace_ndims(M))
 end
-@inline function _array_product_ld(f::F, μ::ProductMeasure, x::AbstractArray, ::IntegerLike) where {F}
+@inline function _array_product_ld(f::F, μ::ProductMeasure, x::AbstractArray, ::Integer) where {F}
     _point_result(_materialize(_batched_kernel(f, μ, x)), μ)
 end
 @inline function _array_product_ld(f::F, μ::ProductMeasure, x::AbstractArray, ::NoMSpaceElementSize) where {F}
@@ -346,7 +346,7 @@ fast_dof(d::AbstractProductMeasure) = _sum_dofs(fast_dof, marginals(d))
 # of freedom each, so their total needs no reduction over the marginals
 # (which may live on a device):
 @inline function _unit_dof(::Type{M}) where {M}
-    static(_static_ndims_of(mspace_ndims(M)) === static(0) && preferred_stdmeasure(M) isa Type{<:StdMeasure})
+    static(mspace_ndims(M) === 0 && preferred_stdmeasure(M) isa Type{<:StdMeasure})
 end
 @inline _sum_dofs(f, mar::StaticArray) = mapreduce(f, +, mar; init = static(0))
 @inline _dynamic_dof(n::IntegerLike) = dynamic(n)
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
index aaa7d0c3..5952a760 100644
--- a/src/combinators/superpose.jl
+++ b/src/combinators/superpose.jl
@@ -177,14 +177,14 @@ end
 # folded pairwise over the component types so that it stays a constant:
 @inline mspace_ndims(::Type{<:SuperpositionMeasure{C}}) where {C<:AbstractArray} = mspace_ndims(eltype(C))
 @inline function mspace_ndims(::Type{MU}) where {C<:Tuple,MU<:SuperpositionMeasure{C}}
-    static_reduce(_CommonNDims{MU}(), mspace_ndims, C)
+    static_mapreduce(mspace_ndims, _CommonNDims{MU}(), C)
 end
 struct _CommonNDims{MU} <: Function end
-@inline (::_CommonNDims{MU})(a::IntegerLike, b::IntegerLike) where {MU} = a == b ? a : NoMSpaceElementSize{MU}()
+@inline (::_CommonNDims{MU})(a::Integer, b::Integer) where {MU} = a == b ? a : NoMSpaceElementSize{MU}()
 @inline (::_CommonNDims{MU})(::Any, ::Any) where {MU} = NoMSpaceElementSize{MU}()
 
 @inline mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:AbstractArray} = _scalar_or_unknown(mspace_flatsize(eltype(C)))
 @inline function mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:Tuple}
-    _scalar_or_unknown(static_reduce(_common_flatsize, mspace_flatsize, C))
+    _scalar_or_unknown(static_mapreduce(mspace_flatsize, _common_flatsize, C))
 end
 @inline _common_flatsize(a, b) = a === b ? a : NoMSpaceElementSize{typeof((a, b))}()
diff --git a/src/density-batched.jl b/src/density-batched.jl
index e50d555d..658e3490 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -61,9 +61,9 @@ end
 
 # The variate rank as a static integer, from the type where known:
 @inline _static_ndims(μ::MU) where {MU} = _static_ndims(mspace_ndims(MU), μ)
-@inline _static_ndims(n::IntegerLike, μ) = static(n)
+@inline _static_ndims(n::Integer, μ) = static(n)
 @inline _static_ndims(::NoMSpaceElementSize, μ) = _static_ndims_of(mspace_ndims(μ))
-@inline _static_ndims_of(n::IntegerLike) = static(n)
+@inline _static_ndims_of(n::Integer) = static(n)
 @inline _static_ndims_of(n::NoMSpaceElementSize) = n
 
 """
@@ -79,7 +79,7 @@ function batched_logdensity_def end
     _default_batched_kernel(logdensity_def, μ, X, _static_ndims(μ))
 end
 
-@inline _default_batched_kernel(f::F, μ, X, n::IntegerLike) where {F} = _default_batched_kernel(f, μ, X, static(n))
+@inline _default_batched_kernel(f::F, μ, X, n::Integer) where {F} = _default_batched_kernel(f, μ, X, static(n))
 @inline _default_batched_kernel(f::F, μ, X, ::StaticInteger{0}) where {F} = _scalar_kernel_broadcast(f, μ, X)
 @inline _default_batched_kernel(f::F, μ, X::AbstractArray, ::StaticInteger{0}) where {F} = _scalar_kernel_broadcast(f, μ, X)
 @inline function _default_batched_kernel(f::F, μ, X::AbstractArray, ::StaticInteger{K}) where {F,K}
@@ -215,7 +215,7 @@ end
 
 # Zero log-densities over the batch dimensions of a flat batch of variates
 # with `n` variate dimensions:
-@inline function _zero_logd_batch(X::AbstractArray, n::IntegerLike)
+@inline function _zero_logd_batch(X::AbstractArray, n::Integer)
     FillArrays.Zeros{_logd_numtype(X)}(ntuple(i -> size(X, n + i), ndims(X) - n))
 end
 @inline _zero_logd_batch(X::AbstractArray{<:Any,N}, ::StaticInteger{N}) where {N} = zero(_logd_numtype(X))
@@ -321,7 +321,7 @@ end
 # in fused operations; otherwise a batch of streams is consumed stream by
 # stream by the outermost stream combinator.
 @inline fixed_stream_size(μ::MU) where {MU} = fixed_stream_size(MU)
-@inline fixed_stream_size(::Type{MU}) where {MU} = static(mspace_ndims(MU) isa IntegerLike)
+@inline fixed_stream_size(::Type{MU}) where {MU} = static(mspace_ndims(MU) isa Integer)
 
 # Batches of streams consumed stream by stream (host loop):
 function _streamwise_ld(f::F, μ, X::AbstractArray) where {F}
diff --git a/src/mspace.jl b/src/mspace.jl
index 01a6f7a8..0ed034e5 100644
--- a/src/mspace.jl
+++ b/src/mspace.jl
@@ -118,10 +118,10 @@ function mspace_ndims end
 @inline mspace_ndims(::Type{MU}) where {MU} = _ndims_of_size(mspace_flatsize(MU), MU)
 @inline mspace_ndims(μ::MU) where {MU} = _ndims_of_size(mspace_flatsize(μ), MU, mspace_ndims(MU))
 
-@inline _ndims_of_size(sz::SizeLike, ::Type) = maybestatic_length(size_dims(sz))
+@inline _ndims_of_size(sz::SizeLike, ::Type) = length(size_dims(sz))
 @inline _ndims_of_size(::NoMSpaceElementSize, ::Type{MU}) where {MU} = NoMSpaceElementSize{MU}()
-@inline _ndims_of_size(sz::SizeLike, ::Type, ::Any) = maybestatic_length(size_dims(sz))
+@inline _ndims_of_size(sz::SizeLike, ::Type, ::Any) = length(size_dims(sz))
 @inline _ndims_of_size(::NoMSpaceElementSize, ::Type, n) = n
 
-@inline _add_ndims(n::IntegerLike, k::IntegerLike) = n + k
-@inline _add_ndims(n::NoMSpaceElementSize, ::IntegerLike) = n
+@inline _add_ndims(n::Integer, k::Integer) = n + k
+@inline _add_ndims(n::NoMSpaceElementSize, ::Integer) = n

From d16fe81b56c20ba390a1e58f85bfcfeebcd02016 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 21 Sep 2026 17:07:46 +0200
Subject: [PATCH 119/122] Take the dimensions of stream multiplicities

Two stream kernels splatted the multiplicity into a size tuple, which
only works for tuples, not for the `StaticArrays.Size` that `SizeLike`
also allows. They take `size_dims` of it now, like the other kernels do.

Created by generative AI.
---
 src/density-batched.jl   | 2 +-
 src/transport-batched.jl | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/density-batched.jl b/src/density-batched.jl
index 658e3490..a7f0d660 100644
--- a/src/density-batched.jl
+++ b/src/density-batched.jl
@@ -278,7 +278,7 @@ end
 @inline function _batched_consume(X::AbstractArray, vsz::SizeLike, sz::SizeLike)
     dims = _consumed_dims(vsz)
     X_flat, X_rest = _batched_split(X, _chunk_rows(prod(dims), sz))
-    return _reshape_consumed(X_flat, (dims..., sz...)), X_rest
+    return _reshape_consumed(X_flat, (dims..., size_dims(sz)...)), X_rest
 end
 @inline _consumed_dims(::Tuple{}) = (static(1),)
 @inline _consumed_dims(vsz::SizeLike) = size_dims(vsz)
diff --git a/src/transport-batched.jl b/src/transport-batched.jl
index 5878e865..9d0fa2e3 100644
--- a/src/transport-batched.jl
+++ b/src/transport-batched.jl
@@ -159,7 +159,7 @@ end
 @inline _merge_multiplicity(Z::AbstractArray, sz::SizeLike) = merge_leading_dims(Z, static(1) + maybestatic_length(size_dims(sz)))
 @inline _split_multiplicity(Z::AbstractArray, ::Tuple{}, n) = Z
 @inline function _split_multiplicity(Z::AbstractArray, sz::SizeLike, n)
-    maybestatic_reshape(Z, (n, sz..., Base.tail(_batch_dims(Z))...))
+    maybestatic_reshape(Z, (n, size_dims(sz)..., Base.tail(_batch_dims(Z))...))
 end
 
 

From c6e4c9eab180b978711b58873f354abc585ee805 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 21 Sep 2026 17:08:05 +0200
Subject: [PATCH 120/122] Draw variates of tuple products without allocating

`_map` passed the mapped function on unspecialized, so drawing the
marginals of a tuple or named tuple product ran through a dynamic call
and boxed every scalar variate. `map` over a named tuple splats its
values on top of that, mapping over the values avoids it. `rand` of a
product of scalar measures, and of a bind over one, allocates nothing
now.

Created by generative AI.
---
 src/combinators/product.jl | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index 7e202e8a..e70e2f26 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -288,11 +288,16 @@ end
 @inline (k::_PointLogd{F,Nothing})(m, x) where {F} = _point_ld(k.f, m, x)
 
 # TODO: Better `map` support in MappedArrays
-_map(f, args...) = map(f, args...)
-_map(f, x::MappedArrays.ReadonlyMappedArray) = mappedarray(fchain((x.f, f)), x.data)
+# `F` keeps the mapped function specialized, it would be passed on
+# unspecialized otherwise and its results boxed:
+_map(f::F, args...) where {F} = map(f, args...)
+# `map` over a named tuple splats its values and boxes the results,
+# mapping over the values themselves doesn't:
+@inline _map(f::F, nt::NamedTuple{names}) where {F,names} = NamedTuple{names}(map(f, values(nt)))
+_map(f::F, x::MappedArrays.ReadonlyMappedArray) where {F} = mappedarray(fchain((x.f, f)), x.data)
 # `map` over a struct array builds struct arrays of the results, variates
 # of the marginals are wanted as plain arrays:
-_map(f, x::StructArray) = [f(m) for m in x]
+_map(f::F, x::StructArray) where {F} = [f(m) for m in x]
 
 function testvalue(::Type{T}, d::AbstractProductMeasure) where {T}
     _map(m -> testvalue(T, m), marginals(d))

From 4fbb67241452291338d1117e274f81d0523afb0c Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 21 Sep 2026 17:08:37 +0200
Subject: [PATCH 121/122] Clean up after the static variate work

Power kernels name the measure again when a variate isn't an array,
instead of letting StaticThings report a dimension mismatch. The shape
checks of powers take `Dims` consistently, the dead `_split_after` for
tuples with a dynamic split point and the dead `_array_product_kernel`
for dynamic ranks are gone, the `StdMeasure` docstring states that
variates of standard measures are `Real`s, and the redesign notes name
the new multiplicity type.

Created by generative AI.
---
 redesign.md                |  4 ++--
 src/collection_utils.jl    |  1 -
 src/combinators/power.jl   | 18 ++++++++++++------
 src/combinators/product.jl |  3 ---
 src/standard/stdmeasure.jl |  7 +++++++
 5 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/redesign.md b/redesign.md
index 1ec36c3d..71b8cbf8 100644
--- a/redesign.md
+++ b/redesign.md
@@ -56,7 +56,7 @@ are accepted, their flat storage is the tuple of component storages.
 inside such streams are flat vectors consumed with the with-rest
 protocol. Point forms return `(result, x_μ, x_rest)` (binds need the
 consumed variate), batched forms take streams `(rows, batch dims...)`
-and a multiplicity `sz::Dims` of variates per stream and return
+and a multiplicity `sz::SizeLike` of variates per stream and return
 `(result, rest)`. Powers pass their size as multiplicity to their base;
 combined measures and tuple products split rows by their fixed stream
 lengths. `fixed_stream_size(::Type{M})` decides whether a batch of
@@ -89,7 +89,7 @@ allocate their own arrays.
 
 **Random variates.** `rand(ctx::GenContext, μ)` with RNG, precision and
 compute unit (`rand(μ)`, `rand(rng, μ)`, `rand(T, μ)` are wrappers).
-`batched_rand_impl(ctx, μ, sz::Dims)` returns a flat batch, a single
+`batched_rand_impl(ctx, μ, sz::SizeLike)` returns a flat batch, a single
 variate for `sz == ()`; `rand_impl` defaults to it. Defaults draw
 standard variates in bulk on the compute unit and transport them, or
 generate variate by variate without a standard transport.
diff --git a/src/collection_utils.jl b/src/collection_utils.jl
index dbf66acc..9b4f2155 100644
--- a/src/collection_utils.jl
+++ b/src/collection_utils.jl
@@ -26,7 +26,6 @@ _exp_cumsum_log(xs::AbstractVector) = exp.(cumsum(log.(xs)))
 Base.@propagate_inbounds _as_tuple(v::AbstractVector, ::Val{N}) where {N} = Tuple(SVector{N}(v))
 
 
-@inline _split_after(x::Tuple, n) = _split_after(x::Tuple, Val{n}())
 @inline _split_after(x::Tuple, ::Val{N}) where {N} = x[begin:(begin+N-1)], x[(begin+N):end]
 
 @generated function _split_after(x::NamedTuple{names}, ::Val{names_a}) where {names,names_a}
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
index 37bc3e7b..a86a06ea 100644
--- a/src/combinators/power.jl
+++ b/src/combinators/power.jl
@@ -135,8 +135,13 @@ end
     _powered_kernel_impl(f, μ, X, _static_ndims(pwr_base(μ)))
 end
 @inline function _powered_kernel_impl(f::F, μ::PowerMeasure, X, ::Any) where {F}
-    sum_leading_dims(_batched_kernel(f, pwr_base(μ), X), static(length(pwr_axes(μ))))
+    _pwr_sum_dims(_batched_kernel(f, pwr_base(μ), X), static(length(pwr_axes(μ))))
 end
+
+# `sum_leading_dims`, but naming the measure when the variate isn't an array:
+@inline _pwr_sum_dims(ℓ, n::StaticInteger) = sum_leading_dims(ℓ, n)
+@inline _pwr_sum_dims(ℓ::Number, ::StaticInteger{0}) = ℓ
+@noinline _pwr_sum_dims(::Number, ::StaticInteger) = _throw_pwr_variate_not_array()
 # Numeric batches of powers of bases without a variate rank are batches of
 # streams:
 @inline function _powered_kernel_impl(::typeof(logdensityof_impl), μ::PowerMeasure, X::AbstractArray{<:Number}, ::NoMSpaceElementSize)
@@ -160,7 +165,7 @@ _powered_stream_kernel(μ::PowerMeasure, X::AbstractArray, ::False) = _streamwis
     end
     return nothing
 end
-@inline _check_pwr_dims(::AbstractArray, ::NoMSpaceElementSize, ::SizeLike, ::Bool) = nothing
+@inline _check_pwr_dims(::AbstractArray, ::NoMSpaceElementSize, ::Dims, ::Bool) = nothing
 @inline batched_logdensityof_impl(μ::PowerMeasure, X) = _powered_kernel(logdensityof_impl, μ, X)
 @inline batched_logdensity_def(μ::PowerMeasure, X) = _powered_kernel(logdensity_def, μ, X)
 
@@ -184,9 +189,10 @@ function _powered_point_nested(f::F, μ::PowerMeasure, x::AbstractArray, ::NoFla
     ν = pwr_base(μ)
     sum(_PointLogd(f, ν), x; init = zero(_logd_numtype(x)))
 end
-@noinline function _powered_point(::F, ::PowerMeasure, x) where {F}
+@noinline _powered_point(::F, ::PowerMeasure, x) where {F} = _throw_pwr_variate_not_array()
+
+@noinline _throw_pwr_variate_not_array() =
     throw(ArgumentError("Variates of powers of measures must be arrays"))
-end
 
 # Nested variates have the power's shape:
 @inline function _check_pwr_shape(μ::PowerMeasure, x::AbstractArray)
@@ -207,7 +213,7 @@ function batched_logdensityof_with_rest(μ::PowerMeasure, x::AbstractVector, sz:
 end
 function _powered_ld_with_rest(μ::PowerMeasure, X::AbstractArray, sz::SizeLike, ::True)
     ℓ, X_rest = batched_logdensityof_with_rest(pwr_base(μ), X, (size_dims(pwr_size(μ))..., size_dims(sz)...))
-    return sum_leading_dims(ℓ, static(length(pwr_axes(μ)))), X_rest
+    return _pwr_sum_dims(ℓ, static(length(pwr_axes(μ)))), X_rest
 end
 function _powered_ld_with_rest(μ::PowerMeasure, x::AbstractVector, ::Tuple{}, ::False)
     ν = pwr_base(μ)
@@ -275,7 +281,7 @@ end
     return nothing
 end
 @inline _check_pwr_flat(x::AbstractArray, k::StaticInteger, dims::Dims) = _check_pwr_dims(x, k, dims, true)
-@inline _check_pwr_flat(::AbstractArray, ::NoMSpaceElementSize, ::SizeLike) = _throw_size_mismatch()
+@inline _check_pwr_flat(::AbstractArray, ::NoMSpaceElementSize, ::Dims) = _throw_size_mismatch()
 
 checked_arg(μ::PowerMeasure, x::Any) = _throw_size_mismatch()
 
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
index e70e2f26..e6a61be6 100644
--- a/src/combinators/product.jl
+++ b/src/combinators/product.jl
@@ -231,9 +231,6 @@ for (bhead, head) in [(:batched_logdensityof_impl, :logdensityof_impl), (:batche
     end
 end
 
-@inline function _array_product_kernel(f::F, μ::ProductMeasure, X::AbstractArray, k::Integer) where {F}
-    _array_product_kernel(f, μ, X, static(k))
-end
 @inline function _array_product_kernel(f::F, μ::ProductMeasure, X::AbstractArray, ::StaticInteger{0}) where {F}
     mar = marginals(μ)
     _check_flatsize(X, maybestatic_size(mar))
diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl
index 9ee949f6..19110167 100644
--- a/src/standard/stdmeasure.jl
+++ b/src/standard/stdmeasure.jl
@@ -1,3 +1,10 @@
+"""
+    abstract type MeasureBase.StdMeasure <: AbstractMeasure
+
+Supertype of the standard measures that transports pivot on.
+
+Variates of standard measures are `Real` numbers.
+"""
 abstract type StdMeasure <: AbstractMeasure end
 
 StdMeasure(::typeof(rand)) = StdUniform()

From badb95726b8a1d6d598ce556558467f1426b39da Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Mon, 21 Sep 2026 17:08:49 +0200
Subject: [PATCH 122/122] Cover the static stream and batch kernels

Adds static stream multiplicities (as a tuple of static integers and as a
`StaticArrays.Size`) through the batched with-rest transport and density
kernels of static powers and combined measures, the batched transports
and `logdensities` on a static batch with the broadcast hook, and
allocation checks for `rand`. The nested-power test checks the layout
instead of pinning the container type.

Created by generative AI.
---
 test/static_variates.jl | 61 +++++++++++++++++++++++++++++++++++------
 1 file changed, 53 insertions(+), 8 deletions(-)

diff --git a/test/static_variates.jl b/test/static_variates.jl
index eeb2e3f5..453504cd 100644
--- a/test/static_variates.jl
+++ b/test/static_variates.jl
@@ -8,10 +8,13 @@ using Test
 
 using MeasureBase
 using MeasureBase: StdNormal, StdUniform, StdExponential
-using MeasureBase: productmeasure, mbind, weightedmeasure, transport_to, logdensityof
+using MeasureBase: productmeasure, mbind, mcombine, weightedmeasure, transport_to, logdensityof
+using MeasureBase: batched_logdensityof_with_rest, batched_transport_to_std,
+    batched_transport_from_std, batched_transport_to_std_with_rest,
+    batched_transport_from_std_with_rest, _materialize
 using MeasureBase.InverseFunctions: inverse
-using ArraysOfArrays: ArrayOfSimilarArrays, flatview, sliced
-using StaticArrays: SVector, SMatrix
+using ArraysOfArrays: flatview, sliced
+using StaticArrays: SVector, SMatrix, Size
 using Static: static
 
 include("testutils.jl")
@@ -37,10 +40,13 @@ const static_model = mbind(static_kernel, static_primary, merge)
         @test @inferred(rand(StdUniform()^static(2))) isa SVector{2,Float64}
         @test @inferred(rand(StdExponential()^static(4))) isa SVector{4,Float64}
 
+        @test allocations_of(rand, StdNormal()^static(3)) == 0
+
         # Nested powers keep the flat `(base dims..., power dims...)` rule:
         x = @inferred(rand((StdNormal()^static(2))^static(3)))
-        @test x isa ArrayOfSimilarArrays{Float64,1,1,<:SMatrix{2,3}}
         @test flatview(x) isa SMatrix{2,3,Float64}
+        @test length(x) == 3 && all(xi -> xi isa SVector{2,Float64}, x)
+        @test reduce(hcat, x) == flatview(x)
 
         μ = StdNormal()^static(3)
         z = SVector(0.1, 0.2, 0.3)
@@ -76,6 +82,7 @@ const static_model = mbind(static_kernel, static_primary, merge)
         @test allocations_of(f, x) == 0
         @test allocations_of(f_inv, z) == 0
         @test allocations_of(logdensityof, μ, x) == 0
+        @test allocations_of(rand, μ) == 0
 
         # The model density is the sum of the component densities:
         x_a = (a = x.a, b = x.b)
@@ -106,6 +113,8 @@ const static_model = mbind(static_kernel, static_primary, merge)
         @test allocations_of(f, x) == 0
         @test allocations_of(f_inv, z) == 0
         @test allocations_of(logdensityof, μ, x) == 0
+        @test allocations_of(rand, μ) == 0
+        @test allocations_of(rand, productmeasure((a = StdNormal(), b = StdExponential()))) == 0
 
         # Tuple products behave the same way:
         μ_t = productmeasure((StdNormal(), StdUniform()^static(2)))
@@ -118,12 +127,48 @@ const static_model = mbind(static_kernel, static_primary, merge)
 
     @testset "batches of static variates" begin
         μ = StdNormal()^static(3)
+        ν = StdUniform()^static(3)
         X = SMatrix{3,4}(reshape(collect(1:12) ./ 10, 3, 4))
-        ℓ = logdensities(μ, X)
-        @test ℓ ≈ [logdensityof(μ, SVector{3}(X[:, i])) for i in 1:4]
 
-        ν = StdUniform()^static(3)
+        @test @inferred(logdensities(μ, X)) ≈
+              [logdensityof(μ, SVector{3}(X[:, i])) for i in 1:4]
+        @test allocations_of(logdensities, μ, X) == 0
+
+        Z = @inferred batched_transport_to_std(StdUniform, μ, X)
+        @test Z ≈ reduce(hcat, [transport_to(ν, μ)(SVector{3}(X[:, i])) for i in 1:4])
+        @test @inferred(batched_transport_from_std(StdUniform, μ, Z)) ≈ X
+
+        # The broadcast hook transports the whole batch at once:
         Y = transport_to(ν, μ).(sliced(X, Val(1)))
-        @test flatview(Y) ≈ reduce(hcat, [transport_to(ν, μ)(SVector{3}(X[:, i])) for i in 1:4])
+        @test flatview(Y) ≈ Z
+        @test Y[2] ≈ transport_to(ν, μ)(SVector{3}(X[:, 2]))
+    end
+
+    # Several variates per stream, with the multiplicity as a tuple of
+    # static integers and as a `StaticArrays.Size`:
+    @testset "static stream multiplicity" begin
+        μ = StdNormal()^static(2)
+        mc = mcombine(vcat, StdNormal()^static(2), StdUniform()^static(3))
+        x = SVector{4}(randn(4))
+        xc = SVector{10}(vcat(randn(2), rand(3), randn(2), rand(3)))
+        to_u = transport_to(StdUniform(), StdNormal())
+
+        for sz in ((static(2),), Size(2))
+            z, x_rest = batched_transport_to_std_with_rest(StdUniform, μ, x, sz)
+            @test z isa SVector{4,Float64} && isempty(x_rest)
+            @test z ≈ to_u.(x)
+            x_back, z_rest = batched_transport_from_std_with_rest(StdUniform, μ, z, sz)
+            @test x_back ≈ reshape(x, (2, 2)) && size(z_rest, 1) == 0
+
+            ℓ, x_ld_rest = batched_logdensityof_with_rest(μ, x, sz)
+            @test _materialize(ℓ) isa SVector{2,Float64} && isempty(x_ld_rest)
+            @test _materialize(ℓ) ≈ [logdensityof(μ, x[(2i - 1):(2i)]) for i in 1:2]
+
+            zc, xc_rest = batched_transport_to_std_with_rest(StdUniform, mc, xc, sz)
+            @test length(zc) == 10 && isempty(xc_rest)
+            ℓc, xc_ld_rest = batched_logdensityof_with_rest(mc, xc, sz)
+            @test isempty(xc_ld_rest)
+            @test _materialize(ℓc) ≈ [logdensityof(mc, xc[(5i - 4):(5i)]) for i in 1:2]
+        end
     end
 end