User experience: Python vs R, PyMC vs Stan, PyTensor vs JAX

[ricardoV94]

@bob-carpenter

I get the part about (embedded) PPLs, but I thought PyMC was more flexible about output format. Thomas Wiecki sent me a notebook producing JAX and some examples under pymc_experimental using Blackjax, but the latter are now 404. Is that just embedding into JAX rather than translating into JAX? We can embed Stan models in JAX that way, but it doesn’t let us jit them or run them on GPUs

I understood your dissatisfaction was about the input, that you can’t write a PyMC model the way you would write a STAN model, line by line. Not about the output.

PyMC is perfectly capable of outputting to you the model conditional/joint densities in jax code (which you called translated into JAX?) which you can then transform/vmap/pmap as you please. That’s how we interop with the jax samplers (ot rust via the numba backend)