TensorFlow Probability

0.25.0 · active · verified Sun Apr 12

TensorFlow Probability (TFP) is a Python library built on TensorFlow that facilitates probabilistic reasoning and statistical analysis. It seamlessly integrates probabilistic models with deep learning on modern hardware (TPUs, GPUs) by providing a wide selection of probability distributions, bijectors, probabilistic layers, variational inference, Markov chain Monte Carlo, and optimizers. Currently at version 0.25.0, TFP follows a regular release cadence, typically aligning with TensorFlow releases.

Warnings

breaking TensorFlow Probability is NOT compatible with Keras 3. Starting with TensorFlow 2.16+, `pip install tensorflow` installs Keras 3 by default. Attempting to use TFP with Keras 3 will result in `AttributeError` upon import or usage of `tfp.layers`.
Fix: Explicitly install Keras 2 by running `pip install tf-keras` or `pip install tensorflow-probability[tf]`. Alternatively, set the environment variable `TF_USE_LEGACY_KERAS=1` before importing TensorFlow and TFP to force TensorFlow to use Keras 2 if `tf-keras` is installed.
gotcha TFP versions are tested and stable against specific TensorFlow versions. Using incompatible versions can lead to unexpected errors or instability.
Fix: Always check the TFP release notes for the recommended compatible TensorFlow version. For example, TFP 0.25.0 is stable against TensorFlow 2.18.
breaking Support for Python 3.8 was removed in TensorFlow Probability 0.22.0.
Fix: Upgrade your Python environment to 3.9 or newer. TFP 0.24.0 and later support Python 3.12.
gotcha When using Python 3.12, a TensorFlow + `wrapt` bug can cause errors in `tfp.layers` and `tfp.experimental.nn`.
Fix: Set the environment variable `WRAPT_DISABLE_EXTENSIONS=true` to work around the bug.
gotcha TensorFlow Probability on JAX (using `tensorflow_probability.substrates.jax`) has some functional limitations compared to the TensorFlow backend.
Fix: Be aware that `tfp.layers` is not fully supported, `tfp.math.minimize` does not work due to `tf.Variable` dependence, and `tf.Variable` equivalents are not present. Pseudorandomness handling also differs.

Install

pip install tensorflow-probability Stable release
pip install tensorflow-probability[tf] With TensorFlow Keras 2 compatibility
pip install tensorflow-probability[tf-cpu] With TensorFlow CPU Keras 2 compatibility
pip install tensorflow-probability[jax] With JAX substrate

Imports

tfp
```
import tensorflow_probability as tfp
```

tfd

import tensorflow_probability.distributions as tfd

tfb

import tensorflow_probability.bijectors as tfb

tfp.glm
```
import tensorflow_probability.glm as tfg
```
The Generalized Linear Model module is aliased as `glm`, not `general_linear_model`.

Quickstart

This quickstart demonstrates fitting a simple Bayesian logistic regression model using TensorFlow Probability's Generalized Linear Model (GLM) module. It generates synthetic data from a Bernoulli distribution, defines a Bernoulli GLM, and then fits the model to estimate coefficients.

import tensorflow as tf
import tensorflow_probability as tfp

tfd = tfp.distributions

# Generate synthetic data for logistic regression
features = tfd.Normal(loc=0., scale=1., name='features').sample(int(100e3))
labels = tfd.Bernoulli(logits=1.618 * features, name='labels').sample()

# Specify a Bernoulli GLM (Generalized Linear Model)
model = tfp.glm.Bernoulli()

# Fit the model using Maximum Likelihood Estimation
coeffs, linear_response, is_converged, num_iter = tfp.glm.fit(
    model_matrix=features[:, tf.newaxis],
    response=tf.cast(labels, dtype=tf.float32),
    model=model
)

print(f"Estimated Coefficients: {coeffs.numpy()}")
print(f"Converged: {is_converged.numpy()}")

view raw JSON →