{"id":25237080,"url":"https://github.com/lockwo/distreqx","last_synced_at":"2025-06-12T07:32:42.393Z","repository":{"id":235005225,"uuid":"789896382","full_name":"lockwo/distreqx","owner":"lockwo","description":"Distrax, but in equinox. Lightweight JAX library of probability distributions and bijectors. ","archived":false,"fork":false,"pushed_at":"2025-04-16T00:33:41.000Z","size":3999,"stargazers_count":30,"open_issues_count":9,"forks_count":5,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-05-12T17:46:57.968Z","etag":null,"topics":["jax","machine-learning","statistics"],"latest_commit_sha":null,"homepage":"https://lockwo.github.io/distreqx/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lockwo.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":".github/FUNDING.yml","license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null},"funding":{"github":["lockwo"]}},"created_at":"2024-04-21T20:56:52.000Z","updated_at":"2025-05-02T02:08:18.000Z","dependencies_parsed_at":"2024-05-10T21:24:08.616Z","dependency_job_id":"09d7b453-4e4f-40ce-93ef-2e604d878f96","html_url":"https://github.com/lockwo/distreqx","commit_stats":null,"previous_names":["lockwo/distreqx"],"tags_count":2,"template":false,"template_full_name":null,"purl":"pkg:github/lockwo/distreqx","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lockwo%2Fdistreqx","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lockwo%2Fdistreqx/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lockwo%2Fdistreqx/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lockwo%2Fdistreqx/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lockwo","download_url":"https://codeload.github.com/lockwo/distreqx/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lockwo%2Fdistreqx/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":259420581,"owners_count":22854616,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["jax","machine-learning","statistics"],"created_at":"2025-02-11T15:33:27.918Z","updated_at":"2025-06-12T07:32:42.372Z","avatar_url":"https://github.com/lockwo.png","language":"Python","funding_links":["https://github.com/sponsors/lockwo"],"categories":["Libraries"],"sub_categories":[],"readme":"\u003ch1 align='center'\u003edistreqx\u003c/h1\u003e\n\u003ch2 align='center'\u003eDistrax + Equinox = distreqx. Easy Pytree probability distributions and bijectors.\u003c/h2\u003e\n\n[![DOI](https://zenodo.org/badge/789896382.svg)](https://zenodo.org/doi/10.5281/zenodo.13764511)\n\ndistreqx (pronounced \"dist-rex\") is a [JAX](https://github.com/google/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.).\n\nThe origin of this package is a reimplementation of [distrax](https://github.com/google-deepmind/distrax), (which is a subset of [TensorFlow Probability (TFP)](https://github.com/tensorflow/probability), with some new features and emphasis on jax compatibility) using [equinox](https://github.com/patrick-kidger/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.)\n\nCurrent features include:\n\n- Probability distributions\n- Bijectors\n\n\n## Installation\n\n```\npip install distreqx\n```\n\nor\n\n```\ngit clone https://github.com/lockwo/distreqx.git\ncd distreqx\npip install -e .\n```\n\nRequires Python 3.9+, JAX 0.4.11+, and [Equinox](https://github.com/patrick-kidger/equinox) 0.11.0+.\n\n## Documentation\n\nAvailable at https://lockwo.github.io/distreqx/.\n\n## Quick example\n\n```python\nimport jax\nfrom jax import numpy as jnp\nfrom distreqx import distributions\n\nkey = jax.random.PRNGKey(1234)\nmu = jnp.array([-1., 0., 1.])\nsigma = jnp.array([0.1, 0.2, 0.3])\n\ndist = distributions.MultivariateNormalDiag(mu, sigma)\n\nsamples = dist.sample(key)\n\nprint(dist.log_prob(samples))\n```\n\n## Differences with Distrax\n\n- No official support/interoperability with TFP\n- 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](https://docs.kidger.site/equinox/tricks/#ensembling) of a `ScalarAffine`)\n- Broader pytree enablement \n- Strict [abstract/final](https://docs.kidger.site/equinox/pattern/) design pattern\n\n## Citation\n\nIf you found this library useful in academic research, please cite: \n\n```bibtex\n@software{lockwood2024distreqx,\n  title = {distreqx: Distributions and Bijectors in Jax},\n  author = {Owen Lockwood},\n  url = {https://github.com/lockwo/distreqx},\n  doi = {10.5281/zenodo.13764512},\n}\n```\n\n(Also consider starring the project on GitHub.)\n\n## See also: other libraries in the JAX ecosystem\n\n[GPJax](https://github.com/JaxGaussianProcesses/GPJax): Gaussian processes in JAX. \n\n[flowjax](https://github.com/danielward27/flowjax): Normalizing flows in JAX.\n\n[Optimistix](https://github.com/patrick-kidger/optimistix): root finding, minimisation, fixed points, and least squares.  \n\n[Lineax](https://github.com/patrick-kidger/lineax): linear solvers.  \n\n[sympy2jax](https://github.com/patrick-kidger/sympy2jax): SymPy\u003c-\u003eJAX conversion; train symbolic expressions via gradient descent.  \n\n[diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable.\n\n[Awesome JAX](https://github.com/lockwo/awesome-jax): a longer list of other JAX projects.  \n\n## Original distrax copyright\n\n```\nCopyright 2021 DeepMind Technologies Limited. All Rights Reserved.\n\nLicensed under the Apache License, Version 2.0 (the \"License\");\nyou may not use this file except in compliance with the License.\nYou may obtain a copy of the License at\n\n    http://www.apache.org/licenses/LICENSE-2.0\n\nUnless required by applicable law or agreed to in writing, software\ndistributed under the License is distributed on an \"AS IS\" BASIS,\nWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\nSee the License for the specific language governing permissions and\nlimitations under the License.\n==============================================================================\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flockwo%2Fdistreqx","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flockwo%2Fdistreqx","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flockwo%2Fdistreqx/lists"}