Parax: Parametric Modeling in JAX + Equinox [P]
Our take
Hi everyone!
Just wanted to share my Python project Parax - an add-on on top of the Equinox library catering for parameter-first modeling in JAX.
For our scientific applications, we found that we often needed to attach metadata to our parameter objects, such as marking them as fixed or attached a prior probability distribution. Further, we often needed to manipulate these parameters in very deep hierarchies, which sometimes can be unintuitive using eqx.tree_at.
We therefore developed Parax, which providesparax.Parameter and parax.Module (that both inherit from eqx.Module) as well as a few helper utilities. These provide a more object-orientated model inspection and manipulation approach, while still following Equinox's immutable principles.
There is some documentation along with a few examples. Perhaps the package is of use to someone else out there! :)
Cheers,
Gary
[link] [comments]
Read on the original site
Open the publisher's page for the full experience
Related Articles
- Parax v0.5: Parametric Modeling in JAX [P]Hi everyone! Just sharing an update on my project Parax, which caters for "parametric modeling" in JAX. Previously, Parax was more focused on scientific applications, however I've since generalized it to be a tool useful for any type of JAX work. It now has a strong focus on a clean, extandable API, as well as ensuring the library is entirely opt-in, as opposed to its previous versions which took a more framework-like approach. Some of Parax's features: Derived/constrained parameters with metadata Computed PyTrees and callable parameterizations Abstract interfaces for fixed, bounded, and probabilistic PyTrees and parameters Filtering and manipulation tools The documentation is available here along with some basic examples. Perhaps the package is of use to someone out there! Cheers, Gary submitted by /u/gvcallen [link] [comments]
- Parax v0.7: Parametric Modeling in JAX [P]Hi everyone! Parax is a library for "Parametric modeling" in JAX, attempting to bridge the approach between pure JAX PyTrees, and more object-orientated modeling approaches (e.g. using Equinox). v0.7 has been released, featuring a more polished API as well as some detailed examples in the documentation. Some of Parax's features: Derived/constrained parameters with metadata Computed PyTrees and callable parameterizations Abstract interfaces for fixed, bounded, and probabilistic PyTrees and parameters Two new examples in the docs that show off these features Bounded optimization (JAXopt) Bayesian sampling (BlackJAX) Perhaps the library is of use to someone, and feel free to leave any feedback! Cheers, Gary submitted by /u/gvcallen [link] [comments]