---
id: distrax
version: "0.1.9"
license: unclear
license_treatment: permissive
maintenance: active
---

# distrax — Distrax: Probability distributions in JAX.

License: permissive · Maintenance: active · Downloads: 188.5K/mo

## 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 above — 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

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

- License: not declared (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 188.5K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

jax probability distributions, bijectors invertible functions, tensorflow probability alternative, jax random sampling, probabilistic distributions jax, machine learning distributions, reinforcement learning policies, jax-ecosystem, probabilistic-modeling, reinforcement-learning

[View on SkillFed](https://skillfed.io/packages/distrax) · [View on PyPI](https://pypi.org/project/distrax/)