Hi, I would like to use a GPU and Jax but I did only find the pymc.sampling_jax — PyMC dev documentation documentation as source code. Is there a documentation to the individual functions etc. like for everything else?

Do you mean these? pymc.sampling_jax.sample_numpyro_nuts pymc.sampling_jax.sample_blackjax_nuts

Sampling Jax documentation?

Oliver October 5, 2022, 10:12am 3

thanks, exactly thats what I was looking for! thanks!

Topic		Replies	Views
Jax sampling in Pymc4: no pm.sampling.jax module v5	5	1667	April 10, 2024
Confusion between latest release version and development version of PyMC Questions	7	543	October 31, 2021
Clear PyMC workflow for beginners Questions	1	822	March 19, 2022
Sampling_jax issues v5	6	1338	July 10, 2023
Set up environment for JAX sampling with GPU support in PyMC v4 Sharing gpu , sampling , ubuntu	16	17324	August 15, 2023