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?
1 Like
Do you mean these?
pymc.sampling_jax.sample_numpyro_nuts
pymc.sampling_jax.sample_blackjax_nuts
3 Likes
thanks, exactly thats what I was looking for! thanks!