Lightweight JAX library of probability distributions and bijectors.
Project description
distreqx
Distrax + Equinox = distreqx. Easy Pytree probability distributions and bijectors.
distreqx (pronounced "dist-rex") is a JAX-based library providing implementations of distributions, bijectors, and tools for statistical and probabilistic machine learning with all benefits of jax (native GPU/TPU acceleration, differentiability, vectorization, distributing workloads, XLA compilation, etc.).
The origin of this package is a reimplementation of distrax, (which is a subset of TensorFlow Probability (TFP), with some new features and emphasis on jax compatibility) using equinox. As a result, much of the original code/comments/documentation/tests are directly taken or adapted from distrax (original distrax copyright available at end of README.)
Current features include:
- Probability distributions
- Bijectors
Installation
pip install distreqx
or
git clone https://github.com/lockwo/distreqx.git
cd distreqx
pip install -e .
Requires Python 3.9+, JAX 0.4.11+, and Equinox 0.11.0+.
Documentation
Available at https://lockwo.github.io/distreqx/.
Quick example
import jax
from jax import numpy as jnp
from distreqx import distributions
key = jax.random.PRNGKey(1234)
mu = jnp.array([-1., 0., 1.])
sigma = jnp.array([0.1, 0.2, 0.3])
dist = distributions.MultivariateNormalDiag(mu, sigma)
samples = dist.sample(key)
print(dist.log_prob(samples))
Differences with Distrax
- No official support/interoperability with TFP
- The concept of a batch dimension is dropped. If you want to operate on a batch, use
vmap
(note, this can be used in construction as well, e.g. vmaping the construction of aScalarAffine
) - Broader pytree enablement
- Strict abstract/final design pattern
Citation
If you found this library useful in academic research, please cite:
@software{lockwood2024distreqx,
title = {distreqx: Distributions and Bijectors in Jax},
author = {Owen Lockwood},
url = {https://github.com/lockwo/distreqx},
doi = {[tbd]},
}
(Also consider starring the project on GitHub.)
See also: other libraries in the JAX ecosystem
GPJax: Gaussian processes in JAX.
flowjax: Normalizing flows in JAX.
Optimistix: root finding, minimisation, fixed points, and least squares.
Lineax: linear solvers.
sympy2jax: SymPy<->JAX conversion; train symbolic expressions via gradient descent.
diffrax: numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable.
Awesome JAX: a longer list of other JAX projects.
Original distrax copyright
Copyright 2021 DeepMind Technologies Limited. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
File details
Details for the file distreqx-0.0.1.tar.gz
.
File metadata
- Download URL: distreqx-0.0.1.tar.gz
- Upload date:
- Size: 35.0 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/5.1.1 CPython/3.12.6
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | 71a95ff3c349b753d90244cc2cf9c70601fb37ddcacf5681cf5d8ff96c8013c5 |
|
MD5 | 1a37fb0ffdd2e5f3b350f3af025083f0 |
|
BLAKE2b-256 | 1b4012801085131278a8e47cdfb61e4be302ae89219b682f29558e07d938785c |
File details
Details for the file distreqx-0.0.1-py3-none-any.whl
.
File metadata
- Download URL: distreqx-0.0.1-py3-none-any.whl
- Upload date:
- Size: 46.0 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/5.1.1 CPython/3.12.6
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | b684c812445149081482c5364fd744835d577165096b5f342f836a6417741263 |
|
MD5 | 1fde4b6058dcb1d591f26649fede885d |
|
BLAKE2b-256 | 30beacdfd3ba520c9653d6675abca974b6b09022929daa6060c47d87acac1039 |