Hello, I’ve been reading this post ( MCMC for big datasets: faster sampling with JAX and the GPU - PyMC Labs ) on how JAX backend (using numpyro) along with GPU can speed up NUTS sampler. But I was wondering, would that also apply to the variational inference algorithm in Pymc? Or currently we just h…

Does JAX backend Speedup also apply to variational inference?

frankligy January 20, 2023, 3:12pm 3

very appreciate for the information!

Topic		Replies	Views
Jax sampling for bayesian neural network v5 jax , pytensor	3	252	March 27, 2024
Use of numpyro/Jax with pymc-dev v5 jax	5	2359	July 28, 2022
Clear PyMC workflow for beginners Questions	1	822	March 19, 2022
Introducing bayeux Sharing	9	662	April 10, 2024
Native JAX support in PyMC without wrappers?	3	260	March 21, 2024