So, I was really excited to be able to use jax in pymc4 but I am having a hard time setting it up. I followed the instructions in Set up environment for JAX sampling with GPU support in PyMC v4 - #4 by payamphysics Everything seems to be set up. Jax imports and reports jax.default_backend() to be …

Got it! you have to import the module separately: import pymc.sampling.jax as pmjax You get a warning, saying that this module is experimental but then you can do: pmjax.sample_numpyro_nuts(1000)

Jax sampling in Pymc4: no pm.sampling.jax module

Questions v5

joshualeond April 10, 2024, 3:38am 5

Hi! I’m new to the PyMC world. Is this still the recommended way of sampling with JAX? Thanks!

1 Like

Topic		Replies	Views
Sampling_jax issues v5	6	1338	July 10, 2023
Confusion between latest release version and development version of PyMC Questions	7	543	October 31, 2021
Installation guide for pymc3jax v3	6	1614	March 18, 2022
Clear PyMC workflow for beginners Questions	1	822	March 19, 2022
Set up environment for JAX sampling with GPU support in PyMC v4 Sharing gpu , sampling , ubuntu	16	17323	August 15, 2023

Jax sampling in Pymc4: no pm.sampling.jax module

Related topics