try enable Mooncake for the LDA model

[yebai]

The variables z in the LDA model are discrete and treated as model parameters,

DynamicPPL.jl/benchmarks/src/Models.jl

Line 150 in 5c89efc

z[i] ~ Categorical(θ[d[i]])

In principle, differentiating w.r.t discrete parameters z won't work. However, ReverseDiff seems to work fine, thus motivating trying Mooncake in this PR.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 84.56%. Comparing base (0fcca13) to head (411ed44).

Additional details and impacted files

@@                         Coverage Diff                         @@
##           mhauru/dot-tilde-tests-from-turing     #841   +/-   ##
===================================================================
  Coverage                               84.56%   84.56%           
===================================================================
  Files                                      34       34           
  Lines                                    3830     3830           
===================================================================
  Hits                                     3239     3239           
  Misses                                    591      591           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Benchmark Report for Commit 411ed4446c3f63f7c6fffac0eb7f19220000d26e

Computer Information

Julia Version 1.11.4
Commit 8561cc3d68d (2025-03-10 11:36 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × AMD EPYC 7763 64-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver3)
Threads: 1 default, 0 interactive, 1 GC (on 4 virtual cores)

Benchmark Results

|                 Model | Dimension |  AD Backend |      VarInfo Type | Linked | Eval Time / Ref Time | AD Time / Eval Time |
|-----------------------|-----------|-------------|-------------------|--------|----------------------|---------------------|
| Simple assume observe |         1 | forwarddiff |             typed |  false |                  9.4 |                 1.6 |
| Simple assume observe |         1 | reversediff |             typed |  false |                  9.4 |                29.8 |
| Simple assume observe |         1 |    mooncake |             typed |  false |                  9.4 |                 8.7 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                732.4 |                36.0 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                413.5 |                48.3 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |               1213.8 |                28.3 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               3598.0 |                22.3 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1475.0 |                29.5 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |               1477.4 |                 3.5 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               5459.5 |                 5.1 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |               1115.6 |                 8.4 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              60707.8 |                 4.3 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               9056.3 |                 9.6 |
|               Dynamic |        10 |    mooncake |             typed |   true |                125.6 |                11.9 |
|              Submodel |         1 |    mooncake |             typed |   true |                 26.2 |                 7.3 |
|                   LDA |         6 |    mooncake |             typed |   true |                106.4 |                 8.5 |
|                   LDA |         6 | reversediff |             typed |   true |                106.9 |                13.5 |

[yebai]

@willtebbutt @mhauru, please note the relatively poor performance of Mooncake on the LDA model, likely caused by a lack of compintell/Mooncake.jl#508. It is also interesting to note that ReverseDiff, though based on a worse approach than Mooncake, has better performance thanks to its linearised tape.

[willtebbutt]

What's the evidence that the lack of linearisation causes this performance problem? Mine and @mhauru 's experience has been that, whenever Mooncake is substantially slower than ReverseDiff, it's because the original code is type unstable.

[willtebbutt]

Yeah, a cursory look at this example suggests that it's super type-unstable -- there's Vector{Vector{Real}} everywhere. I would be willing to bet that fixing that will resolve the problem.

edit: recall that while Mooncake will run on type-unstable code, we make no performance promises.

[yebai]

What's the evidence that the lack of linearisation causes this performance problem? Mine and @mhauru 's experience has been that, whenever Mooncake is substantially slower than ReverseDiff, it's because the original code is type unstable.

I overlooked the type instability issue.

[willtebbutt]

Ahh cool. Looks like the timings make more sense now?

[mhauru]

Thanks for picking this up, I noted this as well and meant to come back to it.