jaxmod 0.2.4

Lightweight utility functions for JAX arrays, batching, and pytrees

Jaxmod

About

Jaxmod—short for "JAX for modelling"—is a lightweight "extension pack" for JAX-based scientific computing. It provides convenience utilities, wrappers, and conventions on top of the excellent Equinox and Optimistix libraries, making them even easier to use for real scientific workflows.

The library was originally created to avoid code duplication in thermochemistry applications for planetary science. Many such problems share similar structural patterns—stoichiometric systems, batched equilibrium calculations, constraints, and differentiable solvers. Jaxmod consolidates this boilerplate into reusable components and establishes consistent conventions, whether for vectorising models efficiently or extending solver behaviour in a principled way.

Documentation

The documentation is available online, with options to download it in EPUB or PDF format:

https://jaxmod.readthedocs.io/en/latest/

Quick install

Jaxmod is a Python package that can be installed on a variety of platforms (e.g. Mac, Windows, Linux). It is recommended to install Jaxmod in a dedicated Python environment. Before installation, create and activate the environment, then run:

pip install jaxmod