Distrax

Common errors

• TypeError: Invalid type for distribution parameter. Expected `jax.Array` or a type convertible to `jax.Array`, but got `float`.

  cause Passing a standard Python float or integer directly to a distribution parameter instead of a JAX array.
  fix Convert parameters to `jax.numpy.array` explicitly, e.g., `loc = jnp.array(0.0)`.

• ValueError: sample requires a PRNG key.

  cause Calling `distribution.sample()` or `sample_and_log_prob()` without providing a JAX PRNG key via the `seed` argument.
  fix Generate 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'

  cause Attempting to use the `BatchReinterpreted` distribution, which was removed in Distrax version 0.1.5.
  fix Refactor code to use `distribution.batch_shape.transpose_event_axes` or `distribution.batch_shape.expand_event_dims` to manipulate batch and event dimensions.

Warnings

• breaking The `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.
  Fix: Ensure `shift` and `scale` are explicitly `jax.numpy.array` scalars or shaped arrays, rather than Python floats, if you encounter dimension-related errors.
• deprecated The `BatchReinterpreted` distribution was deprecated in version 0.1.5 and subsequently removed. Attempting to use it in newer versions will result in an `AttributeError`.
  Fix: Refactor code to use `distribution.batch_shape.transpose_event_axes` or `distribution.batch_shape.expand_event_dims` for manipulating batch and event dimensions, as appropriate.
• gotcha All 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.
  Fix: Always generate a JAX PRNG key (e.g., `key = jax.random.PRNGKey(0)`) and pass it as `seed=key` to sampling methods.
• gotcha Distribution 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.
  Fix: Convert distribution parameters to JAX arrays explicitly, e.g., `loc=jnp.array(0.0)`.

Install

• pip install distrax Install distrax

Imports

• 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

Quickstart

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}")