Rogue PiNodes: PiNode(:(Distributions.nothing), Nothing)

[penelopeysm]

Now that Mooncake works on 1.11 (🎉🎉) Turing's CI reveals some more exciting test failures 😄

Specifically, truncated distributions seem to consistently throw an error. If it helps, this only happens when a distribution is truncated on one side: truncated(Beta(2, 2); lower=0.1, upper=0.9) runs fine.

using DifferentiationInterface
import Mooncake
using Distributions: Beta, truncated, logpdf

function fn(x)
    return logpdf(truncated(Beta(2, 2); lower=0.1), x)
end
backend = AutoMooncake(; config=nothing)
prep = prepare_gradient(fn, backend, [0.5])

The environment contains [email protected], [email protected], [email protected].

Traceback

ERROR: Mooncake.UnhandledLanguageFeatureException("PiNode: Core.PiNode(:(Distributions.nothing), Nothing)")
Stacktrace:
  [1] unhandled_feature(msg::String)
    @ Mooncake ~/ppl/Mooncake.jl/src/interpreter/ir_utils.jl:260
  [2] make_ad_stmts!(stmt::Core.PiNode, line::Mooncake.ID, info::Mooncake.ADInfo)
    @ Mooncake ~/ppl/Mooncake.jl/src/interpreter/s2s_reverse_mode_ad.jl:392
  [3] _broadcast_getindex_evalf
    @ ./broadcast.jl:673 [inlined]
  [4] _broadcast_getindex
    @ ./broadcast.jl:646 [inlined]
  [5] getindex
    @ ./broadcast.jl:605 [inlined]
  [6] macro expansion
    @ ./broadcast.jl:968 [inlined]
  [7] macro expansion
    @ ./simdloop.jl:77 [inlined]
  [8] copyto!
    @ ./broadcast.jl:967 [inlined]
  [9] copyto!
    @ ./broadcast.jl:920 [inlined]
 [10] copy
    @ ./broadcast.jl:892 [inlined]
 [11] materialize(bc::Base.Broadcast.Broadcasted{…})
    @ Base.Broadcast ./broadcast.jl:867
 [12] (::Mooncake.var"#212#214"{Mooncake.ADInfo})(primal_blk::Mooncake.BBlock)
    @ Mooncake ~/ppl/Mooncake.jl/src/interpreter/s2s_reverse_mode_ad.jl:906
 [13] iterate
    @ ./generator.jl:48 [inlined]
 [14] _collect(c::Vector{…}, itr::Base.Generator{…}, ::Base.EltypeUnknown, isz::Base.HasShape{…})
    @ Base ./array.jl:800
 [15] collect_similar
    @ ./array.jl:709 [inlined]
 [16] map
    @ ./abstractarray.jl:3371 [inlined]
 [17] build_rrule(interp::Mooncake.MooncakeInterpreter{…}, sig_or_mi::Type; debug_mode::Bool, silence_debug_messages::Bool)
    @ Mooncake ~/ppl/Mooncake.jl/src/interpreter/s2s_reverse_mode_ad.jl:903
 [18] build_rrule
    @ ~/ppl/Mooncake.jl/src/interpreter/s2s_reverse_mode_ad.jl:856 [inlined]
 [19] prepare_pullback(::Function, ::AutoMooncake{Nothing}, ::Vector{Float64}, ::Tuple{Bool})
    @ DifferentiationInterfaceMooncakeExt ~/.julia/packages/DifferentiationInterface/tBGja/ext/DifferentiationInterfaceMooncakeExt/onearg.jl:13
 [20] prepare_gradient(::typeof(fn), ::AutoMooncake{Nothing}, ::Vector{Float64})
    @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/tBGja/src/first_order/gradient.jl:70
 [21] top-level scope
    @ REPL[75]:1

[penelopeysm]

There's also this PiNode that popped up in a separate test:

Mooncake.UnhandledLanguageFeatureException("PiNode: Core.PiNode(:(Base._msk64), Core.Const(0xffffffffffffffff))")

The full traceback is here: https://github.com/TuringLang/Turing.jl/actions/runs/11449590718/job/31855476593?pr=2341#step:8:1010

This one is harder to narrow down to a proper MWE because it seems to specifically only error with a Gibbs sampler, but in the meantime, I'm triggering the error by checking out TuringLang/Turing.jl#2341 and running this.

using Turing
import Mooncake
obs = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1]
@model function g(obs)
    p ~ Beta(2, 2)
    x ~ Bernoulli(p)
    for i in 1:length(obs)
        obs[i] ~ Bernoulli(p)
    end
    return p, x
end
gibbs = Gibbs(HMC(0.2, 3, :p; adtype=AutoMooncake(;config=nothing)), PG(10, :x))
sample(g(obs), gibbs, 10)

# works just fine
# sample(g(obs), HMC(0.2, 3; adtype=AutoMooncake(;config=nothing)), 10)

[willtebbutt]

Aha! This is very fun. I very much appreciate these examples that you've managed to find. The appearance of constants in PiNodes wasn't really a thing in 1.10, but it appears that in 1.11 the compiler has decided that it's something it wants to do. I'll add proper support for this today!

[willtebbutt]

Okay. This should be resolved in 0.4.19 -- @penelopeysm could you please confirm if this is the case, and close this issue if so?

[penelopeysm]

I think they're all resolved! Lots of green ticks in CI :)