numba backend: compile memory is superlinear in wrapper size

[bridgeland]

I stumbled upon a bug in nutpie: conditioning a LogNormal consumes far more memory than conditioning a Gamma, and the ratio scales with model size.

Note that while I encountered the bug, I leaned on Claude Opus both to diagnose and to craft the bug report. If Claude made mistakes, mea culpa.

Description below; code to reproduce attached.

nutpie_repro.py

---

numba backend: compile memory is superlinear in wrapper size — a LogNormal likelihood costs 3–8× a Gamma one, while PyTensor compiles both identically

Summary

Two models that differ only in which likelihood expresses the same mean and
coefficient of variation compile at essentially identical cost under PyTensor,
but through nutpie.compile_pymc_model(backend="numba") the LogNormal version
needs several times more memory than the Gamma version — and the ratio grows
with model size. On a real model this reached 62 GB, where the Gamma variant
needed 7.4 GB, and the process was killed by the OS.

The likelihoods are equivalent parameterizations:

# gamma
pm.Gamma(name, mu=mu, sigma=mu * cv, observed=y)

# lognormal — same mean, same coefficient of variation
sigma = pt.sqrt(pt.log(1.0 + cv**2))
pm.LogNormal(name, mu=pt.log(mu) - 0.5 * sigma**2, sigma=sigma, observed=y)

Reproduction

nutpie_repro.py (attached) is self-contained — no data files, only
numpy/pymc/pytensor/nutpie. It builds a pool of scalar parameters that every
group depends on, plus a per-group scale parameter, and compiles each variant
in a fresh subprocess with fresh PyTensor and numba caches.

python nutpie_repro.py

      how       kind  groups  shared  peak_rss_GB   wall_s
 pytensor      gamma       6       8         0.25      2.5
 pytensor  lognormal       6       8         0.24      2.3
 pytensor      gamma      12      16         0.28      4.0
 pytensor  lognormal      12      16         0.27      3.5
 pytensor      gamma      18      24         0.32      6.0
 pytensor  lognormal      18      24         0.29      5.0
   nutpie      gamma       6       8         0.56     11.5
   nutpie  lognormal       6       8         0.73     11.5
   nutpie      gamma      12      16         0.84     19.2
   nutpie  lognormal      12      16         1.97     23.3
   nutpie      gamma      18      24         1.21     29.5
   nutpie  lognormal      18      24         4.13     45.0

Under PyTensor both stay at ~0.3 GB and the LogNormal is marginally cheaper.
Under nutpie the LogNormal/Gamma ratio goes 1.3× → 2.3× → 3.4× as the model
grows. Extrapolated to the real model it is 8.4×.

Note for macOS reproducers: RSS badly understates this because of the memory
compressor — the real model reported 6 GB RSS at a 62 GB physical footprint.
Use vmmap --summary <pid> | grep "Physical footprint (peak)".

What this is not

Measured on the real model (400 opportunities, ~47 free RVs), physical
footprint:

what	peak
pytensor.function([logp, *grads]), numba linker, fast_run	823 MB
same, with the fusion rewrite excluded	555 MB
nutpie's own rewrite-then-grad sequence (rewrite_graph(logp, include=["canonicalize", "stabilize"]) then pt.grad)	418 MB
the expand graph (301 outputs)	11 MB
nutpie.compile_pymc_model	62.1 GB (killed)

So it is not the likelihood, not the gradient, not the rewrites, not the
deterministics, and not elemwise fusion. Every component compiles in under a
gigabyte; only the full nutpie path explodes.

Also ruled out:

• Large constants. The whole logp+grad graph carries 206 KB of constants,
  largest 19 KB — so the "promote large constants to shared variables"
  workaround discussed in Make the numba wrapper functions disk-cacheable (and faster to call) #334 does not apply here.
• cache=True. Forwardable via compile_pymc_model(**kwargs) →
  numba.cfunc(c_sig, **kwargs). It buys compile time on a second process, not
  memory, and every model in a batch is a fresh cache key.

Suspected mechanism

_compile_pymc_model_numba wraps the compiled logp and expand functions in
numba.cfunc(c_sig, **kwargs), which compiles the whole chain into a single C
callback. LLVM cost is superlinear in function size, so a modest increase in
the fused graph produces a large increase in compile memory.

Inspecting the fused graph of the real model at equal scale, the difference is
concentrated in one Composite — and notably it is a scalar one, fusing the
gradients of ~28 scalar parameters across the whole model, not anything to do
with the per-observation vectors:

	gamma	lognormal
largest Composite, inner scalar ops	248	344
its inputs / outputs	207 / 28	279 / 28
TrueDiv ops inside	4	58

The LogNormal's per-group sigma = sqrt(log(1 + cv**2)) and the /(2*sigma**2)
term in its logp add gradient structure for each scalar cv, growing that one
Composite by ~39%. That 39% becomes 3–8× in compile memory.

Related

• Make the numba wrapper functions disk-cacheable (and faster to call) #334 (closed unmerged) describes the same wrapper-inlining behaviour: "because
  the wrapper inlines everything up to the dispatcher call, that compilation
  traverses the entire model graph... costs minutes of compile time and a
  GB-scale arena that numba never releases." It was withdrawn because its
  benchmark compared caching-off against caching-on, but that critique does not
  touch the superlinearity shown here: both columns above are cold, and both
  use the same cache settings.
• Slow initialization running PyMC model #115 — slow initialization, closed with "I don't think we can do much about
  the numba compilation times in nutpie".
• High Memory Usage During Sampling #265 — high memory during sampling (different phase).

Environment

nutpie 0.16.10 | pymc 6.0.1 | pytensor 3.0.5
numba 0.65.1 | llvmlite 0.47.0 | numpy 2.4.6
python 3.14.5 | Darwin 25.6.0 arm64

nutpie 0.16.11 contains no compile-path changes, so this is expected to be
unchanged on latest.

Practical impact

For a model of this shape the numba backend is simply unusable with the
LogNormal likelihood: it needs ~62 GB to compile something PyTensor compiles in
under a gigabyte. The alternatives are all closed off — the C backend cannot
compile the graph at all (NotImplementedError in
pytensor/tensor/subtensor.py c_code for an advanced-indexing op), and the
JAX backend is a separate dependency with its own op coverage gaps.

[jessegrabowski]

Can you share exactly what the nutpie_repro.py script is? I didn't see it attached, would be useful for starting digging.

jk i found it, sorry

[velochy]

So, at least my AI friend believes that while the issue is not quite the same I ran into, the same fixes work.
To be more precise, it roughly seems that pytensor #2361 (or numba #10782) fixes the memory blowup and #2354 fixes the compile time regression.

Would be nice if you could try present pytensor head with those two PRs merged in and see if you can confirm.

Just a heads up: I'm just another person that ran head first into the same issues and spent some time babysitting an AI for these PRs. The main maintainer (@ricardoV94) is on the case, but it seems the PRs touch some pretty sensitive code so there is likely to be some back and forth before things actually get merged.

[bridgeland]

Thanks @velochy . I will try that and report back.

("my AI friend" — I am stealing that phrase for my own use.)

[bridgeland]

So, at least my AI friend believes that while the issue is not quite the same I ran into, the same fixes work. To be more precise, it roughly seems that pytensor pymc-devs/pytensor#2361 (or numba numba/numba#10782) fixes the memory blowup and pymc-devs/pytensor#2354 fixes the compile time regression.

Would be nice if you could try present pytensor head with those two PRs merged in and see if you can confirm.

Confirmed. It worked, both for the repro script, and in my production model.

[ricardoV94]

Closing here as it's not a nutpie responsibilty. Thanks for bringing it up @bridgeland