World age weirdness in codegen

[cscherrer]

The way SymPy is set up requires loading some code within an __init__(). Here's a simplified version of what I have:

function __init__()
    stats = PyCall.pyimport_conda("sympy.stats", "sympy")
    SymPy.import_from(stats)
    global stats = stats
    
    @eval begin
        function Distributions.Normal(μ::Sym, σ::Sym)
            println("Evaluating Normal(",μ," :: Sym, ", σ, " :: Sym)")
            stats.$dist(:dist, μ,σ) |> SymPy.density
        end

        Distributions.$dist(μ,σ) = $dist(promote(μ,σ)...)
    end    


end

For a given model, symlogpdf is a GG function that uses SymPy to find its logpdf. Somewhat surprisingly, it works just fine:

julia> m = @model begin
           x ~ Normal()
       end;

julia> truth = rand(m());

julia> symlogpdf(m)
Evaluating Normal(0.0 :: Sym, 1.00000000000000 :: Sym)
       2                       log(π)   log(2)
- 0.5⋅x  - 0.693147180559945 - ────── + ──────
                                 2        2   

The next GG function, codegen, builds code that includes a call to symlogpdf. The code isn't anything too unusual:

function codegen(m::JointDistribution,x)
    return _codegen(m.model, m.args, x)    
end

@gg function _codegen(_m::Model, _args, _data)  
    type2model(_m) |> sourceCodegen() |> loadvals(_args, _data)
end

export sourceCodegen
function sourceCodegen()
    function(m::Model)
        body = @q begin end

        for (x, rhs) in pairs(m.vals)
            push!(body.args, :($x = $rhs))
        end

        push!(body.args, codegen(symlogpdf(m)))
        return body
    end
end

There's a codegen(::Sym) method that does the heavy lifting. This breaks:

julia> codegen(m(),truth)
ERROR: MethodError: no method matching Normal(::SymPy.Sym, ::SymPy.Sym)
The applicable method may be too new: running in world age 26973, while current world is 27213.

It's very strange that the first works but the second doesn't.

@thautwarm I'm pretty stuck on this one, do you have any idea what might be going on here? All relevant code is in master.

[cscherrer]

I changed codegen around completely, things are working better now:

julia> m
@model begin
        σ ~ Normal()
        β ~ Normal()
        α ~ Normal()
        x ~ Normal() |> iid(10000)
        h = α .+ β .* x
        y ~ For(10000) do j
                Normal(h[j], abs(σ))
            end
    end
julia> truth = rand(m());
julia> @btime logpdf($m(),$truth)
  145.049 μs (12 allocations: 78.50 KiB)
-17568.14404327815
julia> @btime $(Soss._codegen(m, false))(NamedTuple(),$truth)
  219.184 μs (2 allocations: 78.20 KiB)
-17568.14404327815
julia> @btime $(Soss._codegen(m, true))(NamedTuple(),$truth)
  19.754 μs (2 allocations: 78.20 KiB)
-17568.144043278204