Bug in nutpie+numba: memory needed for LogNormal compilation scales superlinearly

I discovered a bug in nutpie compilation. A model that conditions a LogNormal consumes far more memory in nutpie compilation than an otherwise identical model that conditions a Gamma. The ratio of memory usage between LogNormal and Gamma scales with model size. On my production model, the LogNormal version exceeds 60Gb while the Gamma version uses less than 8Gb.

Environment: nutpie 0.16.10, pymc 6.0.1, pytensor 3.0.5, numba 0.65.1

See github issue for details. (Is it appropriate to post here given I have already opened an issue on github?)

Definitely good to post here too. I think @Velochy and @ricardoV94 were working on patches for exactly this issue. You try installing the latest pymc 3.3.1 and pytensor from main to see if it moves the needle, or we should be doing a release pretty soon here if you don’t want to do all the package dancing that implies.

1 Like

Actually, my memory problems blew up with pytensor 3.1.0 onward (and 3.0.4, very briefly), so 3.0.5 hints at this likely being a separate issue. To be sure, can you try pytensor 3.0.7, which was the last well-behaving version for me.

1 Like

The latest pymc (6.3.1) and pytensor (3.3.0) does not move the needle. Same memory blowup for LogNormal.

1 Like

Tried pytensor v3.0.7. Same results