Order-agnostic dispatch

Let’s say I have a function like:

function f(x::TypeA, y::TypeB, z::TypeC)
   ...
end

For this particular function f it does not matter what the order of the arguments is, meaning f(x::TypeA, y::TypeB, z::TypeC) == f(::TypeB, y::TypeC, z::TypeA) == f(::TypeC, y::TypeA, z::TypeB) etc.

How can I make f work on every combination of arguments without having to write out every single one by hand? Dispatching on a single abstract type would not work because f could have multiple methods, each operating on different types.

f(args...) = _f(_canonical_order(args...)...)

_f(x::TypeA, y::TypeB, z::TypeC) = ...

if _canonical_order can be implemented recursively or in some other clever way.

Can you give a concrete description of the function you are trying to implement? A lot of the examples where argument-order doesn’t matter fall into one of two categories:

  1. Commutative numerical operations on numerical arguments (like + or hypot). In this case you either promote them to a common type with f(args...) = f(promote(args)...) or (in cases where mixed types can be handled more efficiently, like real + complex) you canonicalize the order.
  2. Cases where you should probably use keyword arguments instead.

This particular example came up when working on Copulas.jl.

I have to compute a correction factor, based on two input marginal distributions. The generic case requires a quadrature and a bisection, which is expensive. But there are closed form solutions for a number of special cases.

Essentially:

# closed-form problems
function _nataf_problem(::Distributions.Normal, Fⱼ::Distributions.LogNormal,
                        ρ::Real, nodes::Integer)
    # r(ρ₀) = ρ₀ s/√(exp(s²) - 1)
    ...
end
function _nataf_problem(::Distributions.Uniform, ::Distributions.Normal,
                        ρ::Real, nodes::Integer)
    # r(ρ₀) = ρ₀ √(3/π).
    ...
end
... list goes on with more special cases

But the order of the marginals doesn’t matter, so you also need the reverse:

# now reserve the order and forward
function _nataf_problem(Fᵢ::Distributions.LogNormal, Fⱼ::Distributions.Normal,
                        ρ::Real, nodes::Integer)
    return _nataf_problem(Fⱼ, Fᵢ, ρ, nodes)
end
function _nataf_problem(Fᵢ::Distributions.Normal, Fⱼ::Distributions.Uniform,
                        ρ::Real, nodes::Integer)
    return _nataf_problem(Fⱼ, Fᵢ, ρ, nodes)
end
... etc

For two arguments, this is of course doable to write out. But for more arguments it quickly explodes. So I was thinking, is there a general way to handle this, even if you had more arguments? I could not find one easily…

Positional arguments are fundamentally distinguished by order, and there’s no stable static sort for types to get around this. Two not-great options off the top of my head:

  1. Runtime sort a vector of distributions by x -> objectid(typeof(x)). That’s an internal detail and can easily change across versions or perhaps even processes, so you wouldn’t want this to be static. Then again, you wouldn’t want this at runtime if you want this to be optimized.
  2. @eval loop for each correction factor formula over permutations of an n-tuple of marginals annotation symbols. Yes, this metaprogramming blows up to factorial(n) methods like the manually forwarding methods, but at least you’re only writing out the template.
julia> begin
       using Combinatorics: permutations
       struct A end
       struct B end
       for (x,y) in permutations((:A, :B))
         @eval f(::$x, ::$y) = "AB"
       end
       f(A(),B()), f(B(),A())
       end
("AB", "AB")

Thanks. Yes, this is fighting positional dispatch which is why it’s so awkward to handle. I was somewhat hoping someone else ran into this problem maybe years ago already and came up with a generic solution that I just missed. It seems broad enough that it would come up more often, but alas?

For the two-argument case I did find this solution using Core.applicable:

function nataf_problem(Fᵢ, Fⱼ, ρ, nodes)
    if applicable(_nataf_closed, Fᵢ, Fⱼ, ρ, nodes)
        _nataf_closed(Fᵢ, Fⱼ, ρ, nodes)
    elseif applicable(_nataf_closed, Fⱼ, Fᵢ, ρ, nodes)
        _nataf_closed(Fⱼ, Fᵢ, ρ, nodes)
    else
        _nataf_generic(Fᵢ, Fⱼ, ρ, nodes)  
    end
end

_nataf_closed(::Normal, Fⱼ::LogNormal, ρ, nodes) = ...
_nataf_closed(::Uniform, ::Normal, ρ, nodes)     = ...

This should compile to true/false so it doesn’t cost anything in runtime.

Sort them by their bare typename, no parameters

name = Base.typename(typeof(<your type>)).wrapper

Use a macro or eval to create the function.

Right, that’s the hard part I am trying to figure out.

Maybe something like this?

@generated function _canonical_order(args...)
    names = [Symbol(t) for t in args]
    perm = sortperm(names)
    quote
        args[$perm]
    end
end

# order of a, b, c does not matter, d is always last
f(a, b, c, d) = f(_canonical_order(a, b, c)..., d)

# We can still dispatch on the types of a, b, c
_f(a::Float64, b::Int, c::String, d) = "f(Float, Int, String)"
_f(a::Float64, b::Int, c::Symbol, d) = "f(Float, Int, Symbol)"
julia> # Float64 < Int < String
       _canonical_order(42, 1.0, "foo")
(1.0, 42, "foo")

julia> _canonical_order("foo", 1.0, 42)
(1.0, 42, "foo")

julia> f(42, 1.0, "foo", π)
"f(Float, Int, String)"

julia> f(:foo, 42, 1.0,  π)
"f(Float, Int, Symbol)"

julia> f(:foo, "bar", 1.0, 42)
ERROR: MethodError: no method matching _f(::Float64, ::String, ::Symbol, ::Int64)

EDIT: this is not type-stable as it stands :confused:

this will recurse forever?

this will also sort on type parameters

maybe this is cleaner:

@generated function _canonical_order(args...)
    perm = sortperm([nameof(t) for t in args])
    return :(($(map(i -> :(args[$i]), perm)...),))
end

f(a, b, c, d) = _f(_canonical_order(a, b, c)..., d)

But note that now for every method you add to f you have to write the arguments in the canonical order as well, otherwise you get a method error as in your example.

Yep, that was the point of the distinction between f and _f but I somehow messed up the copy-pasting from my REPL to discourse.

Good catch; it is also type stable.

For future ref, here is an updated version of my code above

@generated function _canonical_order(args...)
    perm = sortperm([nameof(t) for t in args])
    Expr(:tuple, (:(args[$i]) for i in perm)...)
end

# order of a, b, c does not matter, d is always last
f(a, b, c, d) = _f(_canonical_order(a, b, c)..., d)

# We can still dispatch on the types of a, b, c
_f(a::Float64, b::Int, c::String, d) = "f(Float, Int, String)"
_f(a::Float64, b::Int, c::Symbol, d) = "f(Float, Int, Symbol)"


using Test

@test f(42, 1.0, "foo", π) == "f(Float, Int, String)"
@test f(:foo, 42, 1.0,  π) == "f(Float, Int, Symbol)"

@test_throws MethodError f(:foo, "bar", 1.0, 42)

@inferred f(42, 1.0, "foo", π)

Yes, is that a problem?

Yeah, type inference and compilation won’t leverage an Array even with “statically” known values, so we need the generated function to piece together the tuple expression from the permutation.

You could use _canonical_order or a helper to compute the canonical permutation (and I’m pretty sure string sorts only happen one way), then @eval the proper methods. Of course, you still have to write the argument types in some permutation to begin with, and you still risk accidentally writing 2 permutations that specify the same canonical one. However, you can also check if a method already exists before the @eval, and package precompilation catches such method overwriting for you.

Another “limitation” of computing a canonical permutation in @generated functions (or at runtime) is the _f methods’ annotations are practically constrained to exactly match the sorted concrete types. For example, a (::Int, ::Real) method wouldn’t work because while some inputs like (::Int, ::Rational{Int}) work, others like (::Int, ::Float64) will be sorted to the “wrong” (::Float64, ::Int) and fail to dispatch. However, there wasn’t such a thing as a “consistent” sort across the type hierarchy to begin with, and abstract type annotations wouldn’t have worked anyway for metaprogramming all permutations like the 2nd option I mentioned e.g. (::Int, ::Real) and (::Real, ::Int) runs into the classic multiple dispatch ambiguity because positional argument order is that important.

You can do better than every single one by having the dispatch do some of the sorting work for you. E.g. with four arguments.

f(a::TypeA, b::TypeB, c::TypeC, d::TypeD) = ...
f(a::TypeB, b, c, d) = f(b, a, c, d)
f(a::TypeC, b, c, d) = f(c, b, a, d)
f(a::TypeD, b, c, d) = f(d, b, c, a)
f(a::TypeA, b::TypeC, c, d) = f(a, c, b, d)
f(a::TypeA, b::TypeD, c, d) = f(a, d, c, b)
f(a::TypeA, b::TypeB, c::TypeD, d) = f(a, b, d, c)

You clearly still want to meta-program this but the number of methods will increase quadratically instead of factorially.

Never mind the previous solution. This is much more fun. And linear.

f(a, b, c, d) = f(b, c, d, a)
f(a::TypeA, b, c, d) = f(a, c, d, b)
f(a::TypeA, b::TypeB, c, d) = f(a, b, d, c)

That’s clever! But has a couple of major problems:

  1. If there is no match, then this won’t throw a MethodError but instead rotate forever until you hit a stack overflow.
  2. It’s greedy, and has no backtracking. If I had implementations (A,B,C,D) and (B,X,Y,Z) and pass (B,C,D,A) it would hit (::B, _, _, _) but it would never reach (A,B,C,D)?
  3. You are using the fallback slot, so in my use case above I would always go straight to the quadrature path _nataf(::Distribution, ::Distribution, ρ, nodes) because that is typed and your fallback slot is not.

Always having to write arguments of a new method in the canonical order would be cumbersome, and will be a source of new bugs

Another method with slightly different tradeoffs is to define

argpos(::TypeA) = 1
argpos(::TypeB) = 2
argpos(::TypeC) = 3
function f(x, y, z) 
    t = (x, y, z)
    a, b, c = getindex.((t,), invperm(map(argpos, t)))
    _f(a, b, c)
end

Then define _f with concrete types. That way, at least you avoid stack overflow.
Of course, it’s also possible to avoid that for Gunnar’s nice cyclic permutation solution by defining some Unions (to annotate the previously non-typed arguments).

Apologies if this has already been brought up, but have you considered just using keyword arguments here?

i.e. if you just write

f(; x::TypeA, y::TypeB, z::TypeC) = ...

that somewhat naturally takes care of the ordering for you.

I guess the problem here is that the users have to always associate TypeA with x and so on, but thought I’d mention it at least.

I do like the cleverness here, but the overlapping combinations of input (probability distribution) types should be considered.

  • Forwarding methods-based sorting will interfere across combinations. For example, some combinations’ final permutation can start with TypeB, TypeC, or TypeD instead, and defining their respective versions of the f(a::TypeA, b, c, d) = f(a, c, d, b) method will end up incorrectly freezing the 1st argument of any initial permutation and prevent dispatches to the critical f(a, b, c, d) = f(b, c, d, a).
  • argpos(::TypeB) might be 2 for one combination but 1 in another, a similar interference. That can be taken care of by computing a given combination’s sort from a larger canonical sort across the entire parent module, which could be manually specified as well or a precomputed version of the earlier sorts based on type ID or names.
  • 1 function needs to handle multiple combinations, and keyword arguments alone can’t distinguish multiple methods for them.

Yet another idea: implement insertion sort in type space.

First, we need a function that sorts pairs. For n types, one can just list n(n-1)/2 combinations if that is not excessive, using symmetry:

_pairsort(a::T, b::T) where T = (a, b)
_pairsort(a::TypeA, b::TypeB) = (a, b)
_pairsort(a::TypeA, c::TypeC) = (a, c)
_pairsort(b::TypeB, c::TypeC) = (b, c)
_pairsort(a, b) = (b, a)

If this is excessive, because we have too many types, an alternative implementation can be

_rank(::TypeA) = 1
_rank(::TypeB) = 2
_rank(::TypeC) = 3
_pairsort(a, b) = _rank(a) > _rank(b) ? (b, a) : (a, b)

which requires n methods for n types. All that matters is that this is type stable.

Then a simple insertion sort:

_insert(a) = (a, )
_insert(a, b) = _pairsort(a, b)
function _insert(a, b, c...)
    A, B = _pairsort(a, b)
    (A, _insert(B, c...)...)
end
_sort(sorted) = sorted
function _sort(sorted, a, unsorted...)
    _sort(_insert(a, sorted...), unsorted...)
end
canonical_order(args...) = _sort((), args...)

as it is resolved at compile time anyway. Everything is type stable and quite simple.