https://blog.leokrglv.net/notes/Recent content in Notes on Leo's blogHugoen-usSat, 11 Jul 2026 00:00:00 +0200https://blog.leokrglv.net/notes/jax_sampler_amd/Sat, 11 Jul 2026 00:00:00 +0200https://blog.leokrglv.net/notes/jax_sampler_amd/<p>This is a small note on how to use a GPU-based sampler on linux with an AMD GPU.</p> <p>When using <code>PyMC</code>, it is possible to use external samplers, that increase the sampling speed. For CPU-based samplers, the best-known choice is the Rust-based <code>nutpie</code>.</p> <p>If the underlying model has a high number of levels and/or any type of significant hierarchies, one may want to consider a GPU-based sampler. For that, the common solutions are jax-based <code>numpyro</code> and <code>blackjax</code> samplers. To properly perform them, we need a dedicated GPU. As of today, the most solid ecosystem, even on Linux, is, of course, the Nvidia&rsquo;s CUDA. So the most frictionless way to make the common libraries is using the CUDA infrastructure. I am, currently, an AMD user, so making common libraries work to use GPU&rsquo;s capability, is often not trivial.</p>