Skip to main content

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 a ScalarAffine)
  • 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


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

distreqx-0.0.1.tar.gz (35.0 kB view details)

Uploaded Source

Built Distribution

distreqx-0.0.1-py3-none-any.whl (46.0 kB view details)

Uploaded Python 3

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

Hashes for distreqx-0.0.1.tar.gz
Algorithm Hash digest
SHA256 71a95ff3c349b753d90244cc2cf9c70601fb37ddcacf5681cf5d8ff96c8013c5
MD5 1a37fb0ffdd2e5f3b350f3af025083f0
BLAKE2b-256 1b4012801085131278a8e47cdfb61e4be302ae89219b682f29558e07d938785c

See more details on using hashes here.

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

Hashes for distreqx-0.0.1-py3-none-any.whl
Algorithm Hash digest
SHA256 b684c812445149081482c5364fd744835d577165096b5f342f836a6417741263
MD5 1fde4b6058dcb1d591f26649fede885d
BLAKE2b-256 30beacdfd3ba520c9653d6675abca974b6b09022929daa6060c47d87acac1039

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page