paramax 0.0.4

Parameterizations and parameter constraints for JAX PyTrees.

Paramax

Parameterizations and constraints for JAX PyTrees

Paramax allows applying custom constraints or behaviors to PyTree components, using unwrappable placeholders. This can be used for

• Enforcing positivity (e.g., scale parameters)
• Structured matrices (triangular, symmetric, etc.)
• Applying tricks like weight normalization
• Marking components as non-trainable

Some benefits of the unwrappable pattern:

• It allows parameterizations to be computed once for a model (e.g. at the top of the loss function).
• It is flexible, e.g. allowing custom parameterizations to be applied to PyTrees from external libraries
• It is concise

If you found the package useful, please consider giving it a star on github, and if you create AbstractUnwrappables that may be of interest to others, a pull request would be much appreciated!

Documentation

Documentation available here.

Installation

pip install paramax

Example

>>> import paramax
>>> import jax.numpy as jnp
>>> scale = paramax.Parameterize(jnp.exp, jnp.log(jnp.ones(3)))  # Enforce positivity
>>> paramax.unwrap(("abc", 1, scale))
('abc', 1, Array([1., 1., 1.], dtype=float32))

Alternative parameterization patterns

Using properties to access parameterized model components is common but has drawbacks:

• Parameterizations are tied to class definition, limiting flexibility e.g. this cannot be used on PyTrees from external libraries
• It can become verbose with many parameters
• It often leads to repeatedly computing the parameterization

Related

• We make use of the Equinox package, to register the PyTrees used in the package
• This package spawned out of a need for a simple method to apply parameter constraints in the distributions package flowjax