Scaling with Gpu acceleration

[aaelony]

Hello!

Have you considered integrating or collaborating with general-mcmc?

It uses the burn library for backend-agnostic computation (CPU/GPU), which could be a game-changer for scaling. I’m particularly interested in seeing formula-based modeling (like bambi and brms with multilevel, splines, and gp capabilities) that can handle ML-sized datasets (~50k rows). Curious to hear your thoughts on whether CPU or GPU acceleration is of interest for the nuts-rs roadmap.

[aseyboldt]

You can use burn with nuts-rs as it is. If you want to run the whole sampler in burn, you can also implement the Math trait and use burn in all array operations. You can also use the gpu if you use nuts-rs through from its python frontent and use jax for the likelihood.

[aaelony]

Thank-you, using the Math trait is an interesting idea. I need to learn the existing nuts-rs code and also probably wait until burn 0.22 lands.