TensorFlow Distributions

TL;DR

TensorFlow Distributions is a library that provides flexible, differentiable probabilistic building blocks for deep learning, enabling complex high-dimensional models and fast inference.

Contribution

It introduces a modular framework with distributions and bijectors for probabilistic computation within TensorFlow, supporting advanced models like autoregressive flows.

Findings

Supports high-dimensional probabilistic models

Enables fast, stable sampling and statistics computation

Facilitates deep probabilistic programming and inference

Abstract

The TensorFlow Distributions library implements a vision of probability theory adapted to the modern deep-learning paradigm of end-to-end differentiable computation. Building on two basic abstractions, it offers flexible building blocks for probabilistic computation. Distributions provide fast, numerically stable methods for generating samples and computing statistics, e.g., log density. Bijectors provide composable volume-tracking transformations with automatic caching. Together these enable modular construction of high dimensional distributions and transformations not possible with previous libraries (e.g., pixelCNNs, autoregressive flows, and reversible residual networks). They are the workhorse behind deep probabilistic programming systems like Edward and empower fast black-box inference in probabilistic models built on deep-network components. TensorFlow Distributions has proven an important part of the TensorFlow toolkit within Google and in the broader deep learning community.