klax 0.2.0

A lightweight machine learning package for computational mechanics.

A lightweight machine learning package for computational mechanics built on JAX.

---

Check out the Documentation for examples and reference material.

What is Klax?

Klax provides:

• Specialized machine learning architectures: MLPs with customizable initialization, fully and partially input convex neural networks (ICNNs), matrix-valued neural networks, e.g., skew symmetric matrices, and more.
• Parameter constraints: Differentiable and non-differentiable parameter constraints, e.g., non-negativity and symmetry constraints.
• Highly customizable training and logging utlities: Methods for calibrating abitrary trainable PyTrees with custom loss functions, callbacks, and metrics logging.
• Full JAX compatibility: Seamless integration with JAX's automatic differentiation and acceleration

Klax is build around the highly successfull JAX, Equinox, and Optax projects and designed to be minimally intrusive. All models inherit directly from equinox.Module without additional abstraction layers, ensuring full compatibility with the ecosystem.

The constraint system is derived from Paramax's paramax.AbstractUnwrappable, extending it to support non-differentiable/zero-gradient parameter constraints such as ReLU-based non-negativity constraints.

The training utilities (klax.fit, klax.Loss, klax.Callback) are designed to operate on arbitrarily shaped model and data PyTrees, fully utilizing the flexibility of JAX and Equinox. While they cover most common machine learning use cases, as well as our specialized requirements, they remain entirely optional. The meachine learning architectures implemented in Klax work seamlessly in any JAX-compatible training loop.

Currently Klax's training utilities are built around Optax, but different optimization libraries could be supported in the future if desired.

Support us!

If you like using Klax, feel free to leave a GitHub star, and if there is a machine learning architecture or anything else that you think should be included in Klax, please consider opening a PR.

Installation

Klax can be installed via pip using

pip install klax

If you want to add the latest release to your Python uv project run

uv add klax

or directly install the main branch via

uv add "klax @ git+https://github.com/Drenderer/klax.git@main"