Supporting functions with keyword arguments

[penelopeysm]

This is a slight modification of the basic example in the docs:

using Libtask
@noinline function g(t; n="a", p=3.0) # prevent the compiler being too smart
    produce(t)
end
function f()
    for t in 1:2
        g(t; n="b", p=2.0)
        t += 1
    end
    return nothing
end

t = TapedTask(nothing, f)
consume(t)

I've been mucking around with trying to get this to work (in the interests of making Turing models with kwargs work with SMC/PG). TuringLang/Turing.jl#2007 Right now with Libtask@0.9.4, these are the might_produce overloads that I need:

Libtask.might_produce(::Type{<:Tuple{typeof(Core.kwcall),<:NamedTuple,typeof(g),Vararg}}) = true
Libtask.might_produce(::Type{<:Tuple{<:Function,String,Float64,typeof(g),Vararg}}) = true

What this means for Turing

The first overload (with kwcall) seems easy enough to handle from a Turing point of view. The second one is more annoying though. It arises because Julia does this fancy thing of creating a new function and if you print x.args[1] here

Libtask.jl/src/copyable_task.jl

Line 403 in f154425

Meta.isexpr(x, :invoke) && return might_produce(x.args[1].specTypes)

you get

x.args[1] = MethodInstance for var"#g#3"(::String, ::Float64, ::typeof(g), ::Int64)

where I think it gensyms a new symbol each time it needs to (e.g. if the function is redefined).

Unfortunately this means that it's quite hard to statically specify the correct might_produce call as we need to know the types of the keyword arguments. (Thankfully, the order of the keyword arguments is only as defined in the original function; it doesn't matter which way round they're given in the invocation.) We unfortunately also can't use Vararg since it can only appear as the last type parameter. We could use a macro to generate it.

On the plus side though it is at least possible although it will require extra work from the user! The scary thing though is that if you aren't careful and don't define the right might_produce methods it will silently fail.

Proof of principle of it working with Turing (nb. it won't work with main: you'll need to use TuringLang/Turing.jl#2660 + TuringLang/AdvancedPS.jl#118)

using Turing, Libtask

@model function m(y; n=0)
    x ~ Normal(n)
    y ~ Normal(x)
end
Libtask.might_produce(::Type{<:Tuple{typeof(Core.kwcall),<:NamedTuple,typeof(m),Vararg}}) = true
Libtask.might_produce(::Type{<:Tuple{<:Function,<:Real,typeof(m),Vararg}}) = true

mean(sample(m(5.0), PG(20), 1000))          # approx 2.5
mean(sample(m(5.0; n=10.0), PG(20), 1000))  # approx 7.5

[penelopeysm]

FYI @mhauru

[mhauru]

Could we define a macro that defines migth_produce methods? Would be handy even without kwargs.

I'm not sure how to do it with kwargs, but could it ignore kwargs for now? So you would call call @might_produce g(::Int) and that would turn into, I guess

Libtask.might_produce(::Type{<:Tuple{typeof(g),Int}}) = true
Libtask.might_produce(::Type{<:Tuple{typeof(Core.kwcall),<:NamedTuple,typeof(m),Int}}) = true
Libtask.might_produce(::Type{<:Tuple{<:Function,<:Any,typeof(g),Int}}) = true
Libtask.might_produce(::Type{<:Tuple{<:Function,<:Any,<:Any,typeof(g),Int}}) = true

? We would probably need to introspect the number of kwargs even if we ignore their types. I assume that's possible from some Method related type.

[penelopeysm]

We would probably need to introspect the number of kwargs even if we ignore their types. I assume that's possible from some Method related type.

I can't find a way to do so:

julia> function f(a; b=1, c="s")
           nothing
       end
f (generic function with 1 method)

julia> only(methods(f)).sig
Tuple{typeof(f), Any}

julia> only(methods(Core.kwcall, (NamedTuple, typeof(f), Vararg))).sig
Tuple{typeof(Core.kwcall), NamedTuple, typeof(f), Any}

But I suppose somebody on Slack might know.

[penelopeysm]

Yay!!!

julia> Base.kwarg_decl(methods(f)[1])
2-element Vector{Symbol}:
 :b
 :c

[penelopeysm]

using Turing, Libtask

@model function m(y; n=0)
    x ~ Normal(n)
    y ~ Normal(x)
end

macro might_produce_kwargs(fn)
    quote
        Libtask.might_produce(::Type{<:Tuple{typeof(Core.kwcall),<:NamedTuple,typeof($(esc(fn))),Vararg}}) = true
        possible_n_kwargs = unique(map(length ∘ Base.kwarg_decl, methods($(esc(fn)))))
        for n in possible_n_kwargs
            # we only need `Any` and not `<:Any` because tuples are covariant
            kwarg_types = fill(Any, n)
            Libtask.might_produce(::Type{<:Tuple{<:Function,kwarg_types...,typeof($(esc(fn))),Vararg}}) = true
        end
    end
end

@might_produce_kwargs(m)

mean(sample(m(5.0), PG(20), 1000))          # approx 2.5
mean(sample(m(5.0; n=10.0), PG(20), 1000))  # approx 7.5

I only really tested with Turing models, but there's nothing in the macro that specifies that it has to be a Turing model, so I think this should work generally.

[mhauru]

Nice work. Wanna make a PR? This is gonna be very handy in general, I always forgot the whole ::Type{<:Tuple{typeof(blahblah... construct when I had to define these might_produce methods, even with no keyword args involved.