probability: a library for probabilistic reasoning and statistical analysis with deep network integration

What it solves

It provides a framework for probabilistic reasoning and statistical analysis, allowing developers to integrate probabilistic methods with deep learning networks. It solves the challenge of performing complex statistical computations and inference at scale by leveraging hardware acceleration (GPUs) and distributed computation.

How it works

The library is organized into a layered architecture:

• Layer 0 (Numerical Operations): Utilizes TensorFlow's numerical operations, specifically the LinearOperator class for efficient matrix-free implementations.
• Layer 1 (Statistical Building Blocks): Provides a vast collection of probability distributions (tfp.distributions) and reversible transformations of random variables called Bijectors (tfp.bijectors).
• Layer 2 (Model Building): Enables the creation of joint distributions over interdependent variables and probabilistic neural network layers that incorporate uncertainty.
• Layer 3 (Probabilistic Inference): Implements algorithms for approximating integrals, including Markov chain Monte Carlo (MCMC) for sampling, Variational Inference (VI) for optimization, and specialized stochastic optimizers.

Additionally, the library can operate as "Tensor-friendly Probability" when used with JAX.

Who it’s for

It is designed for researchers and developers working on probabilistic machine learning, Bayesian neural networks, and complex statistical modeling who need the scalability of TensorFlow or JAX.

Highlights

• Hardware Acceleration: Supports GPUs and distributed computation for large-scale models.
• Flexible Substrates: Works with both TensorFlow and JAX.
• Comprehensive Toolset: Includes everything from basic distributions to advanced MCMC and Variational Inference algorithms.
• Uncertainty Integration: Offers probabilistic layers to add uncertainty to standard neural network functions.