@measure combinator - macros, TypeVars, and binops

[cscherrer]

I think it's close to working, but...

julia> using StatsFuns

julia> @measure Normal(μ,σ) ≃ (1/sqrt2π) * Lebesgue(X)
ERROR: UndefVarError: X not defined
Stacktrace:
 [1] (::Type{Lebesgue{X}})() at /home/chad/git/Measures.jl/src/basemeasures/lebesgue.jl:4
 [2] Lebesgue(::TypeVar) at /home/chad/git/Measures.jl/src/basemeasures/lebesgue.jl:6
 [3] top-level scope at /home/chad/git/Measures.jl/src/macros.jl:41

[1] and [2] here are

struct Lebesgue{X} end
Lebesgue(X) = Lebesgue{X}()

Any suggestions on getting this to work properly? Here's the code the macro currently generates:

julia> using MacroTools

julia> (@macroexpand @measure Normal(μ,σ) ≃ (1/sqrt2π) * Lebesgue(X)) |> MacroTools.prettify
quote
    struct Normal{P, X} <: MeasureTheory.AbstractMeasure{X}
        par::P
    end
    function Normal(nt::NamedTuple)
        P = typeof(nt)
        return Normal{P, eltype(Normal{P})}(nt)
    end
    Normal(; kwargs...) = Normal((; kwargs...))
    (baseMeasure(μ::Normal{P, X}) where {P, X}) = (1 / sqrt2π) * Lebesgue(X)
    Normal(μ, σ) = Normal(; μ, σ)
    ((::Normal{P, X} ≪ ::typeof((1 / sqrt2π) * Lebesgue(X))) where {P, X}) = true
    ::typeof((1 / sqrt2π) * Lebesgue(X)) ≪ ::Normal{P, X} where {P, X} = true
end

[phipsgabler]

(baseMeasure(μ::Normal{P, X}) where {P, X}) = (1 / sqrt2π) * Lebesgue(X)

That line.

Parens confused me...

Why is it necessary to include a free variable in the surface form, anyway? Is it supposed to be universally quantified in X?

Also I'd check hygiene once more. This looks a bit suspicious.

[simeonschaub]

The problem is in this line:

((::Normal{P, X} ≪ ::typeof((1 / sqrt2π) * Lebesgue(X))) where {P, X}) = true

X is a TypeVar here, so you can't use it in regular function calls inside the type signature, since those calls in the constructor need to be resolved before the function participates in dispatch.

[phipsgabler]

But you were right, and I wasn't 🤷

[phipsgabler]

I suppose you should introduce an equivalent of

const $(gensym("basemeasuretype(Normal)")) = typeof((1 / sqrt2π) * Lebesgue(X)))

and use that everywhere instead.

[simeonschaub]

But how does it know what X is?

[simeonschaub]

The best way I see for solving this would probably be some promote-like mechanism.

[phipsgabler]

I'd expect

((::Normal{P, X} ≪ ::M) where {P, X, M}) = M <: typeof((1 / sqrt2π) * Lebesgue(X))

to compile to a constant as well. It might even work if we define all three cases of (Normal, M), (M, Normal), and the fallback (M, M) ~> true. But it looks fishy to me wrt. making this transitive.

[cscherrer]

Why is it necessary to include a free variable in the surface form, anyway? Is it supposed to be universally quantified in X?

Yep, that's it. For a while I had

@measure Normal(μ,σ) ≃ Lebesgue

but that assumes the base measure is "primitive", and doesn't give a way to do e.g. scaled base measures.

Also I'd check hygiene once more. This looks a bit suspicious.

It's intentionally unhygienic, gensyms gum up the works in this case.

The problem is in this line:

((::Normal{P, X} ≪ ::typeof((1 / sqrt2π) * Lebesgue(X))) where {P, X}) = true

X is a TypeVar here, so you can't use it in regular function calls inside the type signature, since those calls in the constructor need to be resolved before the function participates in dispatch.

Ah, right. Maybe the syntax should replace X with eltype(Normal{P}). I think I could get this to work, and it may be a simple fix.

I'd expect

((::Normal{P, X} ≪ ::M) where {P, X, M}) = M <: typeof((1 / sqrt2π) * Lebesgue(X))

to compile to a constant as well. It might even work if we define all three cases of (Normal, M), (M, Normal), and the fallback (M, M) ~> true. But it looks fishy to me wrt. making this transitive.

That's really interesting. Seems like it would work? And it just adds one method.

I've been kind of assuming it will be really hard to get things fully transitive. I think I'd rather the user have the potential for a missing method here and there than a consistently large overhead while we generate lots of methods. Assuming it's easy enough to add more methods, anyway. But auto-adding an O(1) number of methods for a given case seems sensible.

[cscherrer]

This seems to work:

using MeasureTheory
using StatsFuns

##########################################
# Macro builds this part

struct Normal{P, X} <: MeasureTheory.AbstractMeasure{X}
    par::P
end

function Normal(nt::NamedTuple)
    P = typeof(nt)
    return Normal{P, eltype(Normal{P})}(nt)
end

Normal(; kwargs...) = Normal((; kwargs...))

(baseMeasure(μ::Normal{P, X}) where {P, X}) = (1 / sqrt2π) * Lebesgue(eltype(Normal{P}))

Normal(μ, σ) = Normal(; μ, σ)

((::Normal{P, X} ≪ ::typeof((1 / sqrt2π) * Lebesgue(eltype(Normal{P})))) where {P, X}) = true
(::typeof((1 / sqrt2π) * Lebesgue(eltype(Normal{P}))) ≪ ::Normal{P, X}) where {P, X} = true

##########################################
# User adds this method

import Base
Base.eltype(::Type{Normal{P}}) where {P} = Real

##########################################
# Trying it out

julia> Normal()
Normal{NamedTuple{(),Tuple{}},Real}(NamedTuple())

julia> baseMeasure(Normal())
MeasureTheory.ScaledMeasure{Float64,Lebesgue{Real},Real}(-0.9189385332046728, Lebesgue{Real}())

julia> Normal(0.1,0.5)
Normal{NamedTuple{(:μ, :σ),Tuple{Float64,Float64}},Real}((μ = 0.1, σ = 0.5))

julia> Normal(μ=0.1, σ=0.5)
Normal{NamedTuple{(:μ, :σ),Tuple{Float64,Float64}},Real}((μ = 0.1, σ = 0.5))

[cscherrer]

Getting close!

julia> using MeasureTheory

julia> using StatsFuns

julia> @measure Normal(μ,σ) ≃ (1/sqrt2π) * Lebesgue(X)
≪ (generic function with 2 methods)

julia> import Base

julia> Base.eltype(::Type{Normal{P}}) where {P} = Real

julia> Normal()
Normal{NamedTuple{(),Tuple{}},Real}(NamedTuple())

julia> baseMeasure(Normal())
MeasureTheory.ScaledMeasure{Float64,Lebesgue{Real},Real}(-0.9189385332046728, Lebesgue{Real}())

julia> Normal(0.1, 0.5)
Normal{NamedTuple{(:μ, :σ),Tuple{Float64,Float64}},Real}((μ = 0.1, σ = 0.5))

julia> Normal(μ=0.1, σ=0.5)
Normal{NamedTuple{(:μ, :σ),Tuple{Float64,Float64}},Real}((μ = 0.1, σ = 0.5))

but ≪ methods aren't right yet:

julia> Normal() ≪  baseMeasure(Normal())
ERROR: MethodError: no method matching ≪(::Normal{NamedTuple{(),Tuple{}},Real}, ::MeasureTheory.ScaledMeasure{Float64,Lebesgue{Real},Real})
Closest candidates are:
  ≪(::Normal{P,X}, ::MeasureTheory.ScaledMeasure{Float64,Lebesgue{Any},Any}) where {P, X} at /home/chad/git/Measures.jl/src/macros.jl:61
Stacktrace:
 [1] top-level scope at REPL[10]:1

Here's what it's generating:

julia> using MacroTools

julia> (@macroexpand @measure Normal(μ,σ) ≃ (1/sqrt2π) * Lebesgue(X)) |> MacroTools.prettify
quote
    struct Normal{P, X} <: MeasureTheory.AbstractMeasure{X}
        par::P
    end
    function Normal(nt::NamedTuple)
        P = typeof(nt)
        return Normal{P, eltype(Normal{P})}(nt)
    end
    Normal(; kwargs...) = Normal((; kwargs...))
    (baseMeasure(μ::Normal{P, X}) where {P, X}) = (1 / sqrt2π) * Lebesgue(eltype(Normal{P}))
    Normal(μ, σ) = Normal(; μ, σ)
    ((::Normal{P, X} ≪ ::typeof((1 / sqrt2π) * Lebesgue(eltype(Normal{P})))) where {P, X}) = true
    ((::typeof((1 / sqrt2π) * Lebesgue(eltype(Normal{P}))) ≪ ::Normal{P, X}) where {P, X}) = true
end

resulting in

julia> methods(≪)
# 2 methods for generic function "≪":
[1] ≪(::Normal{P,X}, ::MeasureTheory.ScaledMeasure{Float64,Lebesgue{Any},Any}) where {P, X} in Main at /home/chad/git/Measures.jl/src/macros.jl:61
[2] ≪(::MeasureTheory.ScaledMeasure{Float64,Lebesgue{Any},Any}, ::Normal{P,X}) where {P, X} in Main at /home/chad/git/Measures.jl/src/macros.jl:65

[phipsgabler]

I'm baffled that this is accepted, but I fear it doesn't work as expected:

≪(::Normal{P,X}, ::MeasureTheory.ScaledMeasure{Float64,Lebesgue{Any},Any})

There's too many Anys there. Compare:

julia> g(x::X, ::typeof(Lebesgue(eltype(Vector{X})))) where {X<:Real} = ()
g (generic function with 1 method)

julia> methods(g)
# 1 method for generic function "g":
[1] g(x::X, ::Lebesgue{Any}) where X<:Real in Main at REPL[70]:1

julia> g(1, Lebesgue(Int))
ERROR: MethodError: no method matching g(::Int64, ::Lebesgue{Int64})
Closest candidates are:
  g(::X, ::Lebesgue{Any}) where X<:Real at REPL[70]:1
Stacktrace:
 [1] top-level scope at REPL[74]:1

julia> g(1, Lebesgue(String))
ERROR: MethodError: no method matching g(::Int64, ::Lebesgue{String})
Closest candidates are:
  g(::X, ::Lebesgue{Any}) where X<:Real at REPL[70]:1
Stacktrace:
 [1] top-level scope at REPL[75]:1

julia> g(1, Lebesgue(Any))
()

probably due to it treating the inner X as a free variable:

julia> XX = TypeVar(:XX)
XX

julia> Vector{XX}
Array{XX,1}

julia> eltype(Vector{XX})
Any

I think this not how Julia is supposed to work; I filed an issue.

[phipsgabler]

Anyway, I find that the syntax is doing too much at the same time, and don't like the free varible. What about this:

# @measure Normal{X} ≃ (1/sqrt2π) * Lebesgue(X)
struct Normal{X, P} <: MeasureTheory.AbstractMeasure{X}
    par::P
end
function Normal(nt::NamedTuple)
    P = typeof(nt)
    return Normal{peltype(Normal, P), P}(nt)
end
basemeasure(::Normal{X}) where {X} = (1/sqrt2π) * Lebesgue(X)
peltype(::Type{Normal}, ::Type) = error("No eltype defined for parametrization with $P")

# @parametrization Normal(μ, σ)::Normal{Real}
Normal(μ, σ) = Normal(; μ, σ)
peltype(::Type{Normal}, ::Type{::NamedTuple{(:μ, :σ), Tuple{Any, Any}}}) = Real

# we can also have complex constraints for some parametrization
# @parametrization Normal(a::T, b::T)::Normal{T} where {T<:Real}
Normal(a::T, b::T) where {T<:Real} = Normal(; a, b)
peltype(::Type{Normal}, ::Type{::NamedTuple{(:a, :b), Tuple{T, T}}}) where {T<:Real} = T

In essence, we have to define, for each measure, a function that gives us the eltype for for each parametrization -- that's what I called peltype for now. And I preferred to have X as the first parameter, since it is more important to the user than the rather internal P, and then we can easily write Normal{Float64} etc.

[mschauer]

Hm, I like to avoid the X in AbstractMeasure{X} by defining eltype functions, but if we have X, isn't X also the eltype?

[cscherrer]

@phipsgabler :

What about this[?]

I think I like it. I mean, it looks good, but there seem to be so many corner cases with this stuff. Let's try it and see how it goes. The Normal{X,P} goes against my intuition, but I see why you prefer it, and I think it might be a good idea.

I notice you leave out the ≪ methods, which may be just as well to start - that seems it can get to be particularly tricky.

[cscherrer]

@mschauer

Hm, I like to avoid the X in AbstractMeasure{X} by defining eltype functions, but if we have X, isn't X also the eltype?

Yes, that's right. But having it statically is important, since then we can dispatch on it. This is a big pain point in Distributions.jl.

[mschauer]

Dispatching on X makes of course a lot of sense, with a hypothetical method improve

improve(::AbstractMeasure{X}) where {X} = ...

It would be nice to have a bit of parallel infrastructure for those who are Measures but cannot inherit AbstractMeasure via vanilla trait dispatch via peltype or eltype. For example for Distributions,

improve(U::D) where {D <: Distribtuion} = improve(U, peltype(D))
improve(U, X) = ...

[phipsgabler]

@cscherrer Yes, I left them out for now, because they seem to be the complicated part. But I think the following is really getting readable:

(mu::Normal{X} ≪ nu::AbstractMeasure) where {X} = typeof(nu) <: typeof((1/sqrt2π) * Lebesgue(X))

Altough I lack intuition whether <: is enough here. Maybe a type equality is needed (l <: r && l >: r). And maybe the macro also should spit out a check that the family of base measures is a concrete type, given the parameters. Like,

let X = Any
    @assert Core.Compiler.isconcretetype(typeof((1/sqrt2π) * Lebesgue(X)))
end

And I was thinking whether a separate type would be appropirate for name-order canonicalization and the design of the trait for sample types:

Normal(; kwargs...) = Normal(Parametrization(; kwargs...))
sampletype(::Normal, p::Parametrization) = error("No sample type defined for parametrization with $(typeof(p))")

# @parametrization Normal(μ, σ)::Normal{Real}
Normal(μ, σ) = Normal(; μ, σ)
sampletype(::Normal, ::Parametrization{(:μ, :σ), Tuple{Any, Any}}) = Real

# @parametrization Normal(; a::T, b::T)::Normal{T} where {T<:Real}
sampletype(::Normal{T}, ::Parametrization{(:a, :b), Tuple{T, T}}) where {T<:Real} = T

where the Parametrization constructor takes care of constructing a named tuple from the kwargs in canonicalized order.

[cscherrer]

@phipsgabler :

But I think the following is really getting readable:

(mu::Normal{X} ≪ nu::AbstractMeasure) where {X} = typeof(nu) <: typeof((1/sqrt2π) * Lebesgue(X))

This does look really nice, but I think it's missing a lot. The = here really needs to be reverse implication. So

typeof(nu) <: typeof((1/sqrt2π) * Lebesgue(X))      implies        (mu::Normal{X} ≪ nu::AbstractMeasure)

In general, ≃ is an equivalence relation, and there seem to be way too many to deal with case-by-case, so we can instead dispatch to some canonical representative of the equivalence class. So for example,

(mu::Normal{X} ≪ nu::AbstractMeasure) where {X} = Lebesgue(X) ≪ nu

I guess we would get to Lebesgue(X) here by starting with mu and recursively call baseMeasure unless we hit a fixpoint. Then we would get a nu with mu ≪ nu. But we have to be careful, maybe stop the recursion when it fails to be an equivalence?

I was thinking whether a separate type would be appropirate for name-order canonicalization and the design of the trait for sample types:

Is a Parameterization here supposed to take the place of P in the original design? What do you see as advantages of doing it this way? Is sampletype(mu) == typeof(rand(mu))? I was thinking this would be the role of X. I think I'm missing some key part of this idea. Could you give some more detail?

[cscherrer]

How about

function representative(μ)
    # Check if we're done
    isprimitive(μ) && return μ

    ν = baseMeasure(μ)
    
    # Make sure not to leave the equivalence class
    ν ≪ μ || return μ

    # Fall back on a recusive call
    return representative(ν)
end


≪(μ, ν) = representative(μ) ≪ representative(ν)

We could have some measures like Lebesgue considered "primitive" (not defined in terms of another measure)

[phipsgabler]

Is a Parameterization here supposed to take the place of P in the original design? What do you see as advantages of doing it this way?

Yes, I thought it might be a good idea to distinguish "canonicalized parametrization" from arbitrary parametrization, allowing certain special operations to be provided by default, and stronger assumptions being made. E.g., one could have

# in the struct, var"#par"::P
const PAR_FIELD = Symbol("#par")
function getproperty(m::Normal{<:Any, <:Parametrization}, name)
    par = getfield(m, PAR_FIELD)
    name == PAR_FIELD ? par : getproperty(par, name)
end

which I wouldn't do for arbitrary types, but only if we know that the parametrization was set up in the special way given through the macro.

And if we do this, I wouldn't assume P always being a NamedTuple, but rather leave that option for the user and rely on a custom type always. A bit like Zygote or DiffRules switched from NamedTuple to Composite for struct adjoints. (For example, I can imagine a measure over named tuples of being parametrized most easily by a named tuple of parameters, which would be precluded or at least more difficult to get right if P<:NamedTuple is "reserved" by the library).

Is sampletype(mu) == typeof(rand(mu))? I was thinking this would be the role of X. I think I'm missing some key part of this idea. Could you give some more detail?

Yes, kind of -- sampletype(M, typeof(p)) >: typeof(rand(M(p)) was the intention. This forms a trait (or functional dependency): for each measure M and parametrization P of the measure, the (super)type of samples X is uniquely determined, isn't it? That's exactly the fact that is used in

function Normal(p::Parametrization)
    P = typeof(p)
    return Normal{sampletype(Normal, P), P}(p)
end

which I forgot to include above. This constructor prevents the user from constructing arbitrary combinations between X and P. After that, we know that for P<:Parametrization the sample type is X == sampletype(m::M{P}) == sampletype(M, P) >: typeof(rand(m)). When the user chooses a custom parametrization, they have to take care themselves that X and P match.

Maybe it should rather be written in trait style, SampleType(M, P) = X. And sampletype(m::M{X, P}) = X as a fallback for P not a Parametrization.

I was just copying the approach as written by you before, because I liked it, but thought that eltype isn't really the right word for this. Sorry if I have cause confusion -- I was just putting forward more ideas. This is such a difficult yet interesting design problem.

[cscherrer]

Thanks @phipsgabler , I think I understand now. One possible consideration in the sampletype/eltype design... In principle, a distribution or measure is a kind of container, and broadcasting over containers is a natural thing to do. Would it make sense to define broadcasting over a measure? Anything this would break, say for infinite support measures?

I agree P shouldn't need to be a named tuple. I have some convenience functions for that case since I think it will be common, but there's no constraint enforcing that. Really it could be any type. What advantage do we gain by adding a <: Parameterization type constraint? And... would this even allow named tuples, since they aren't subtypes of this?

[mschauer]

constructing arbitrary combinations between

In Gaussian the default sample type is the type of the mean, but you can ask for whatever rand(Float32, Normal()) === randn(Float32) making Normal a model for randn(::Type{X})