BART numba optimizations

Hi, I just saw the main pymc release 5.19.0 has new options for numba speed-up with the samplers. I see that the pymc-bart pgbart sampler has numba optimizations written in the code and I am wondering if the new compile_kwargs = dict(mode=“NUMBA”) command will work with BART too or if these optimizations are already being used?

Thanks!

@aloctavodia ?

Hi @JTIEGS, these are two completely separate sets of optimizations. The numba optimizations in PyMC-BART are always used, independent of compile_kwargs = dict(mode="NUMBA") and are exclusively related to the PGBART sampler. The optimizations introduced in PyMC 5.19.0 are more general and should have an overall effect on PyMC, thus they could indirectly also benefit PyMC-BART. I have not tested them yet.

We hope to have a new version of PyMC-BART soon, this new version is being rewritten in Rust and we used better data structures for the trees. So PyMC-BART should become faster and less memory-hungry. Early benchmarks show a speed-up of around 5x or 10x.

1 Like

Great! That’s what I figured, but wanted to double check. Sounds great! I’m looking forward to the new version!

1 Like