Install & Compatibility
Where this runs
tested against v0.1.5 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
muslpy 3.10–3.910 runs
build_error
glibcpy 3.10–3.910 runs
installs and imports cleanly · install 16.0s · import 4.414s · 641MB
665MB installed
● package 665MB
Code
Verified usage
Verified import paths — ran on the pinned version, not inferred.
Categorical
✓ from distrax import Categorical
Normal
✓ from distrax import Normal
Distribution
✓ from distrax import Distribution
Transformed
✓ from distrax import Transformed
Bijector
✓ from distrax import Bijector
This example demonstrates how to define common distributions like Categorical and Normal, sample from them, and compute their log-probabilities using Distrax and JAX. It highlights the requirement for JAX PRNG keys for sampling.
import distrax
import jax
import jax.numpy as jnp
# It's good practice to provide a key for reproducibility
key = jax.random.PRNGKey(0)
# Create a Categorical distribution
probs = jnp.array([0.1, 0.2, 0.7])
categorical = distrax.Categorical(probs=probs)
# Sample from it (requires a JAX PRNG key)
sample = categorical.sample(seed=key)
print(f"Categorical sample: {sample}")
# Compute log-probability
log_prob = categorical.log_prob(sample)
print(f"Categorical log-prob: {log_prob}")
# Create a Normal distribution
loc = jnp.array(0.0)
scale = jnp.array(1.0)
normal = distrax.Normal(loc=loc, scale=scale)
# Sample from it (requires a JAX PRNG key, can specify sample_shape)
sample_normal = normal.sample(seed=key, sample_shape=(5,))
print(f"Normal samples: {sample_normal}")
# Compute log-probability for a specific value
log_prob_normal = normal.log_prob(jnp.array(0.5))
print(f"Normal log-prob of 0.5: {log_prob_normal}")
Debug
Known issues
breakingThe `ScalarAffine` bijector in version 0.1.3 changed its expectation for `shift` and `scale` parameters. They now explicitly expect `Array`s that are scalars or broadcast correctly to event dimensions.fixEnsure `shift` and `scale` are explicitly `jax.numpy.array` scalars or shaped arrays, rather than Python floats, if you encounter dimension-related errors.
affects: >=0.1.3
deprecatedThe `BatchReinterpreted` distribution was deprecated in version 0.1.5 and subsequently removed. Attempting to use it in newer versions will result in an `AttributeError`.fixRefactor code to use `distribution.batch_shape.transpose_event_axes` or `distribution.batch_shape.expand_event_dims` for manipulating batch and event dimensions, as appropriate.
affects: >=0.1.5 (removal)
gotchaAll sampling methods (`sample`, `sample_and_log_prob`) require a JAX PRNG key to be passed via the `seed` argument. Forgetting this will raise an error.fixAlways generate a JAX PRNG key (e.g., `key = jax.random.PRNGKey(0)`) and pass it as `seed=key` to sampling methods.
affects: All versions
gotchaDistribution parameters (e.g., `loc`, `scale`, `probs`) should ideally be JAX arrays (`jax.numpy.array`). Passing standard Python floats or integers might sometimes work due to JAX's auto-conversion, but explicit conversion is recommended to prevent `TypeError` or unexpected broadcasting issues.fixConvert distribution parameters to JAX arrays explicitly, e.g., `loc=jnp.array(0.0)`.
affects: All versions
Errors
Common errors & fixes
TypeError: Invalid type for distribution parameter. Expected `jax.Array` or a type convertible to `jax.Array`, but got `float`.
Passing a standard Python float or integer directly to a distribution parameter instead of a JAX array.
fixConvert parameters to `jax.numpy.array` explicitly, e.g., `loc = jnp.array(0.0)`.
ValueError: sample requires a PRNG key.
Calling `distribution.sample()` or `sample_and_log_prob()` without providing a JAX PRNG key via the `seed` argument.
fixGenerate a JAX PRNG key (`key = jax.random.PRNGKey(0)`) and pass it as `seed=key` to the sampling method.
AttributeError: module 'distrax' has no attribute 'BatchReinterpreted'
Attempting to use the `BatchReinterpreted` distribution, which was removed in Distrax version 0.1.5.
fixRefactor code to use `distribution.batch_shape.transpose_event_axes` or `distribution.batch_shape.expand_event_dims` to manipulate batch and event dimensions.
Upgrade
Version history
0.1.9latest on PyPI · released Jun 12, 2026
Audit
Dependencies
jaxrequiredCore numerical backend for distributions and bijectors.
jaxlibrequiredJAX's compiled core, required for JAX operations.
numpyrequiredUsed internally and for array interoperability.