Parax
Parax is a library for parametric modeling in JAX. Features include:
- Composable, array-like variables with metadata (
Constrained,Random,Derived, etc.), - Unwrappable PyTree parameterizations
- Built-in higher-level bijective constraints (via distreqx)
- Abstract interfaces and associated tree manipulation tools
This makes Parax great for:
- Constraints for machine learning
- Bounded optimization for scientific modeling
- Probabilistic modeling and Bayesian inference
- Deep, nested PyTrees
- Combinations of the above
Note that Parax is not a framework, though it can be used to make one. Rather, it is focused on extensibility and interoperability with other JAX libraries (especially Equinox).
Installation
Parax can be installed using pip:
pip install parax
Documentation
Documentation is available here.
Quick example
Parax provides array-like variables that hold metadata and can be parameterized/constrained:
import parax as prx
import jax.numpy as jnp
p1 = prx.Tagged(1.0, metadata={'hello', 'world'})
p2 = prx.Constrained(prx.constraints.Interval(0.0, 10.0), value=8.0)
p2.raw_value, p2.bounds
# Array(1.3862944), (Array(0.0), Array(10.0))
jnp.sin(p1) + (2 * p2)
# Array(16.84147)
You can also apply arbitrary computations to PyTrees and parameters using explicit unwrapping:
pytree = {'a': 1.0, 'b': {'x': 2.0, 'y': prx.Derived(jnp.log, 3.0)}}
wrapped = prx.Apply(jnp.exp, pytree)
prx.unwrap(wrapped)
# {'a': Array(2.7182817),
# 'b': {'x': Array(7.389056),
# 'y': Array(3.0)}}
In the above example, prx.Apply operates on the whole PyTree's array-like nodes, while prx.Derived is an array-like prx.AbstractVariable.
Motivation
Usually, PyTrees are just "dumb" containers. However, it is often desirable to attach some metadata/parameterization to a specific node. This can be done by "unwrapping" the metadata or constraint during model preparation or computation.
Compared to other approaches, this provides a middle ground between purity and rigidity:
- The "purist" approach is using shadow PyTrees i.e. parallel trees that hold the relevant metadata/parameterization. However, these are tedious to define for nested models, and require the entire library to manage parallel structures.
- The "standard" approach is using properties and attributes i.e. defining the metadata/parameterization implicitly within the model. This is straight-forward, but tightly couples the extra state with the model, resulting in unnecessary fields and computations.
Next steps
Several more involved examples are available in the documentation, for example on bounded optimization and Bayesian sampling.
Related
The library's design was inspired by several others that deserve mention, including Flax, paramax, and PyTorch.
1 maintainer
Release files for parax 0.11.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| parax-0.11.0.tar.gz | 490.1 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| parax-0.11.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 555.3 kB
Release files / parax-0.11.0.tar.gz
| Download URL | parax-0.11.0.tar.gz |
|---|---|
| Size | 490.1 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
3fbe2ab86a1dba66343ade30113e446e4966b65b5b798a27a8d630c2e1df3e7b
|
|
BLAKE2b-256 checksum How to use checksums |
adcca24ff2ebb287a156f33ee31a77ffb5f2100fc8695c38a53d7754db4d5ae2
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Sep 23, 2026.
Transparency logRelease files / parax-0.11.0-py3-none-any.whl
| Download URL | parax-0.11.0-py3-none-any.whl |
|---|---|
| Size | 65.2 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
fbb4f356879d0027ba249b577e1b4052bdc3f3da006601b4326a19d7586a9e62
|
|
BLAKE2b-256 checksum How to use checksums |
3121b1aa2a5717fd972ce31df8a966694be12b0bd6ce407b7cdd7d0ca87d7ba5
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Sep 23, 2026.
Transparency log
