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

For some built-in constraints and probabilistic features, you may need this distreqx branch:

pip install git+https://github.com/gvcallen/distreqx.git

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.

Project links

Data verified by PyPI on Aug 26, 2026
Data provided by the project maintainers, verified at the time the release was uploaded to PyPI.

Key dates

PyPI data
Data sourced directly from PyPI's database.
  • Released:
Newer release available (0.10.8)

1 maintainer

PyPI data
Data sourced directly from PyPI's database.
Avatar for gvcallen from gravatar.com gvcallen

Credits

Author: Gary Allen

GitHub Statistics

Data verified by PyPI on Aug 26, 2026
The GitHub source repository was provided by the project maintainers and verified by PyPI at the time of upload. Stars, forks, and open issues/PRs are derived from that repository and have not been independently verified.

Provides Extra

docs test

Report project as malware

Download files

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

Source Distribution

parax-0.10.7.tar.gz (478.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

parax-0.10.7-py3-none-any.whl (46.3 kB view details)

Uploaded Python 3

File details

Details for the file parax-0.10.7.tar.gz.

File metadata

  • Download URL: parax-0.10.7.tar.gz
  • Upload date:
  • Size: 478.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for parax-0.10.7.tar.gz
Algorithm Hash digest
SHA256 7e628af782a51958c8dcd0ade03c3f09ae5d8edb3280dd6f86837016ce8a9dcf
MD5 20a02b18e6ad58399f72a6ff5cf41b96
BLAKE2b-256 8fd32902a2dea316eecb663f2d3f917c2815db8aefc8e1e67d409b8d2d08626a

See more details on using hashes here.

Provenance

The following attestation bundles were made for parax-0.10.7.tar.gz:

Publisher: publish.yml on gvcallen/parax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file parax-0.10.7-py3-none-any.whl.

File metadata

  • Download URL: parax-0.10.7-py3-none-any.whl
  • Upload date:
  • Size: 46.3 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for parax-0.10.7-py3-none-any.whl
Algorithm Hash digest
SHA256 ee9c0a2711e57010f8b1db3d3997ab8763e1f1570a5c3ca05ac270b0db800de4
MD5 48eaaca7046f878e122b73d6a430135b
BLAKE2b-256 cbb46ae2f581350fbdf9917d8cb2204800d746d51d957a3aa4c747978d7ade19

See more details on using hashes here.

Provenance

The following attestation bundles were made for parax-0.10.7-py3-none-any.whl:

Publisher: publish.yml on gvcallen/parax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Release history Release notifications | RSS feed

0.10.8

2 files

This release

0.10.7 This release

2 files

0.10.6

2 files

0.10.5

2 files

0.10.3

2 files

0.10.2

2 files

0.10.1

2 files

0.10.0

2 files

0.9.9

2 files

0.9.8

2 files

0.9.7

2 files

0.9.6

2 files

0.9.5

2 files

0.9.4

2 files

0.9.3

2 files

0.9.2

2 files

0.9.1

2 files

0.9.0

2 files

0.8.3

2 files

0.8.2

2 files

0.8.0

2 files

0.7.11

2 files

0.7.10

2 files

0.7.9

2 files

0.7.8

2 files

0.7.7

2 files

0.7.6

2 files

0.7.5

2 files

0.7.4

2 files

0.7.3

2 files

0.7.2

2 files

0.7.1

2 files

0.7.0

2 files

0.6.4

2 files

0.6.3

2 files

0.6.2

2 files

0.6.1

2 files

0.6.0

2 files

0.5.12

2 files

0.5.11

2 files

0.5.9

2 files

0.5.8

2 files

0.5.7

2 files

0.5.6

2 files

0.5.4

2 files

0.5.3

2 files

0.5.2

2 files

0.5.0

2 files

0.4.15

2 files

0.4.13

2 files

0.4.12

2 files

0.4.11

2 files

0.4.10

2 files

0.4.9

2 files

0.4.7

2 files

0.4.6

2 files

0.4.4

2 files

0.4.3

2 files

0.4.2

2 files

0.4.1

2 files

0.4.0

2 files

0.3.9

2 files

0.3.8

2 files

0.3.6

2 files

0.3.5

2 files

0.3.4

2 files

0.3.3

2 files

0.2.1

2 files

0.2.0

2 files

0.1.0

2 files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page