Posts tagged JAX

Faster Sampling with JAX and Numba

PyMC can compile its models to various execution backends through PyTensor, including:

Read more ...


Pathfinder Variational Inference

Pathfinder [Zhang et al., 2021] is a variational inference algorithm that produces samples from the posterior of a Bayesian model. It compares favorably to the widely used ADVI algorithm. On large problems, it should scale better than most MCMC algorithms, including dynamic HMC (i.e. NUTS), at the cost of a more biased estimate of the posterior. For details on the algorithm, see the arxiv preprint.

Read more ...


How to wrap a JAX function for use in PyMC

This notebook uses libraries that are not PyMC dependencies and therefore need to be installed specifically to run this notebook. Open the dropdown below for extra guidance.

Read more ...