{"categories":[{"label":"Python Modules","url":"https://skillfed.io/packages/category/software-development-libraries-python-modules/16"},{"label":"Artificial Intelligence","url":"https://skillfed.io/packages/category/scientific-engineering-artificial-intelligence/8"}],"enrichment":{"capability":"Distrax provides JAX-native probability distributions and bijectors (invertible functions with Jacobian computation) as a lightweight reimplementation of a subset of TensorFlow Probability, designed for extensibility and readability.","skillfed_tags":["jax-ecosystem","probabilistic-modeling","reinforcement-learning"],"use_cases":["Build probabilistic agent policies in reinforcement learning with custom distribution definitions.","Sample from and compute log-probabilities of multivariate distributions in JAX-based machine learning pipelines.","Compose complex distributions using bijectors (e.g., transformed distributions via Tanh or other invertible functions).","Migrate TensorFlow Probability code to JAX by using Distrax distributions with cross-compatible APIs.","Implement variational inference or other probabilistic inference methods using JAX's autodiff and JIT compilation."],"what_it_does":"Distrax is a JAX-native probability library that reimplements a subset of TensorFlow Probability with emphasis on readability and extensibility. It provides distributions (like MultivariateNormalDiag) and bijectors (invertible transformations with Jacobian tracking) that work seamlessly with JAX's functional and JIT-compilation paradigms. The library is designed to be cross-compatible with TensorFlow Probability\u2014you can mix Distrax and TFP distributions in the same computation, or wrap one for use in the other's meta-distributions.\n\nThe package targets researchers and practitioners building probabilistic models, particularly in reinforcement learning where custom policy distributions are common. It emphasizes mathematical clarity in implementations and makes it simple to define custom distributions or bijectors. While not intended to replace TensorFlow Probability entirely, it fills the gap for teams already committed to JAX and wanting a lighter-weight, more extensible alternative for probability operations.","worth_installing":"Yes. Distrax is actively maintained, has no known vulnerabilities, and offers a well-designed JAX-native alternative to TensorFlow Probability for teams already using JAX. The permissive Apache 2.0 license and cross-compatibility with TFP make it low-risk to adopt. Install it if you need probability distributions in JAX and value extensibility and mathematical clarity over a comprehensive feature set."},"id":"distrax","links":{"html":"https://skillfed.io/packages/distrax","md":"https://skillfed.io/packages/distrax.md","pypi":"https://pypi.org/project/distrax/"},"maintenance":{"status":"active"},"meta":{"latest_release":"2026-06-12","license_spdx":null,"license_treatment":"permissive","name":"distrax","python_support":"supports_current","summary":"Distrax: Probability distributions in JAX."},"popularity":{"monthly_downloads":188525,"position":9941,"tier":"top_15000"},"security":{"n_vulnerabilities":0},"version":"0.1.9"}