TensorFlow Probability on JAX
https://www.tensorflow.org/probability/examples/TensorFlow_Probability_on_JAX
Run in Google Colab
View source on GitHub
Download notebook
TensorFlow Probability (TFP) is a library for probabilistic reasoning and statistical analysis that now works on JAX! For those not familiar, JAX is a library for accelerated numerical computing based on composable function transformations.
We have ported a lot of TFP's most useful functionality to JAX while preserving the abstractions and APIs that many TFP users are now comfortable with.
