parax 0.5.12

Parametric modeling in JAX

Parax

Parax is a library for parametric modeling in JAX. Features include:

• Derived/constrained parameters with metadata
• Computed PyTrees and callable parameterizations
• Abstract interfaces for fixed, bounded, and probabilistic PyTrees
• Associated filtering and tree manipulation tools

The above makes Parax great for:

• Parameterizations in machine learning
• Bounded optimization in scientific modeling
• Bayesian inference in probabilistic modeling

Examples for each of these are available in the documentation.

Note that Parax is not a framework, and is designed to be both extendable and interoperable with other JAX libraries (especially Equinox).

Installation

Parax can be installed using pip:

pip install parax

For some 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.Param(1.0, metadata={'hello', 'world'})
p2 = prx.Constrained(8.0, prx.Interval(0.0, 10.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 unwrapping:

pytree = {'a': 1.0, 'b': {'x': 2.0, 'y': prx.Derived(3.0, jnp.log)}}
wrapped = prx.Computed(pytree, jnp.exp)

prx.unwrap(wrapped)
# {'a': Array(2.7182817),
#  'b': {'x': Array(7.389056), 
#        'y': Array(3.0)}}

In the above example, prx.Computed operates on the whole PyTree, while prx.Derived is an array-like prx.AbstractVariable.

Related

The library's design was inspired by several others that deserve mention, including Flax, paramax, and PyTorch.