distrax

What it is and what it does

Distrax is a JAX-native probability library that reimplements a subset of TensorFlow Probability with emphasis on readability and extensibility. It provides distributions (like MultivariateNormalDiag) and bijectors (invertible transformations with Jacobian tracking) that work seamlessly with JAX's functional and JIT-compilation paradigms. The library is designed to be cross-compatible with TensorFlow Probability—you can mix Distrax and TFP distributions in the same computation, or wrap one for use in the other's meta-distributions.

The package targets researchers and practitioners building probabilistic models, particularly in reinforcement learning where custom policy distributions are common. It emphasizes mathematical clarity in implementations and makes it simple to define custom distributions or bijectors. While not intended to replace TensorFlow Probability entirely, it fills the gap for teams already committed to JAX and wanting a lighter-weight, more extensible alternative for probability operations.

Use it for:

• Build probabilistic agent policies in reinforcement learning with custom distribution definitions.
• Sample from and compute log-probabilities of multivariate distributions in JAX-based machine learning pipelines.
• Compose complex distributions using bijectors (e.g., transformed distributions via Tanh or other invertible functions).
• Migrate TensorFlow Probability code to JAX by using Distrax distributions with cross-compatible APIs.
• Implement variational inference or other probabilistic inference methods using JAX's autodiff and JIT compilation.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Distrax provides JAX-native probability distributions and bijectors (invertible functions with Jacobian computation) as a lightweight reimplementation of a subset of TensorFlow Probability, designed for extensibility and readability.

Yes. Distrax is actively maintained, has no known vulnerabilities, and offers a well-designed JAX-native alternative to TensorFlow Probability for teams already using JAX. The permissive Apache 2.0 license and cross-compatibility with TFP make it low-risk to adopt. Install it if you need probability distributions in JAX and value extensibility and mathematical clarity over a comprehensive feature set.

Install

distrax on PyPI

pip install distrax

uv add distrax

poetry add distrax

Installing distrax

Before you install

Installation is straightforward with low friction; the package is actively maintained with a recent release and no known vulnerabilities. It depends on JAX, jaxlib, and TensorFlow Probability nightly, which are substantial but standard dependencies in the JAX ecosystem.

License in practice

Licensed under Apache 2.0 (permissive), allowing free use, modification, and distribution in commercial and private projects with minimal restrictions.

Quickstart

pip install distrax

import distrax
import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(1234)
mu = jnp.array([-1., 0., 1.])
sigma = jnp.array([0.1, 0.2, 0.3])
dist = distrax.MultivariateNormalDiag(mu, sigma)
samples = dist.sample(seed=key)
log_prob = dist.log_prob(samples)

Requires Python >= 3.11 and a working JAX installation (including jaxlib, which may require a C++ compiler or prebuilt binaries for your platform).

Verify before relying

• Whether the library's API stability is guaranteed or if breaking changes are expected before a 1.0 release.
• Performance characteristics compared to TensorFlow Probability for common distributions and sampling operations.
• Coverage of which TFP distributions and bijectors are actually implemented in Distrax.

Package facts

Evidence: distrax-0.1.9-py3-none-any.whl

Keywords: jax, probability, distribution, python, machine learning

Tags

More Python Modules packages