Mooncake issue with large NamedTupleDists

[philippeller]

Mooncake AD fails on _flat_ntd_eff_accessors / _flat_ntd_orig_accessors for large NamedTupleDist

Summary

When using Mooncake (reverse-mode AD) with BAT's bat_findmode / OptimizationAlg, the optimization crashes with a MooncakeRuleCompilationError during rule compilation for _flat_ntd_eff_accessors. This happens when the NamedTupleDist prior has many fields (~50 parameters). Smaller priors (~23 parameters) work fine.

Environment

• Julia 1.12.1
• BAT v4.0.4
• ValueShapes v0.11.4
• Mooncake v0.5.23

Minimal reproducer

using BAT, Distributions, ValueShapes, Mooncake, ADTypes

# Build a NamedTupleDist with ~50 fields (mix of Normal, Truncated, Uniform, ConstValueDist)
prior_nt = (;
    (Symbol("param_$i") => Normal(0, 1) for i in 1:30)...,
    (Symbol("trunc_$i") => truncated(Normal(1, 0.1), 0.5, 1.5) for i in 1:10)...,
    (Symbol("unif_$i") => Uniform(0, 1) for i in 1:5)...,
    (Symbol("fixed_$i") => 0.5 for i in 1:5)...,   # these become ConstValueDist
)

prior = distprod(; prior_nt...)
likelihood = let
    params -> begin
        logpdf(Normal(0, 1), sum(values(params))) # dummy
    end
end

posterior = PosteriorMeasure(BAT.logfuncdensity(likelihood), prior)
init = (; (k => rand(v) for (k,v) in pairs(prior))...)

set_batcontext(ad = AutoMooncake(Mooncake.Config()))
bat_findmode(posterior, OptimizationAlg(optalg=Optimization.LBFGS(), init=ExplicitInit([init])))

Error

MooncakeRuleCompilationError: an error occurred while Mooncake was compiling
a rule to differentiate something.

To replicate this error run the following:

Mooncake.build_rrule(
    Mooncake.MooncakeInterpreter(Mooncake.ReverseMode),
    Tuple{
        Type{ValueShapes.NamedTupleShape},
        Type{NamedTuple},
        @NamedTuple{field1::ValueShapes.ArrayShape{Real, 1}, ...50 fields...}
    };
    debug_mode=false
)

The stacktrace points to:

_flat_ntd_eff_accessors
  @ BAT/src/transforms/distribution_transform.jl:658
dist_trafo_impl
  @ BAT/src/transforms/distribution_transform.jl:701-708

[BJMCox]

@oschulz Do we want to support Mooncake given the ongoing situation surrounding it?

[oschulz]

Probably yes, but let's wait until the MeasureBase refactor is done, since a lot of the code that will need specific pullbacks will move there.

[BJMCox]

That makes sense. I will have a look more precisely as to the situation there as well - there have been some rumblings of a replacement, but that is always the case 😆