| Name | reax JSON |
| Version |
0.5.5
JSON |
| download |
| home_page | None |
| Summary | REAX: A simple training framework for JAX-based projects |
| upload_time | 2025-07-20 19:27:50 |
| maintainer | None |
| docs_url | None |
| author | None |
| requires_python | >=3.9 |
| license | None |
| keywords |
machine learning
jax
research
|
| VCS |
 |
| bugtrack_url |
|
| requirements |
No requirements were recorded.
|
| Travis-CI |
No Travis.
|
| coveralls test coverage |
No coveralls.
|
REAX
====
.. image:: https://codecov.io/gh/muhrin/reax/branch/develop/graph/badge.svg
:target: https://codecov.io/gh/muhrin/reax
:alt: Coverage
.. image:: https://github.com/camml-lab/reax/actions/workflows/ci.yml/badge.svg
:target: https://github.com/camml-lab/reax/actions/workflows/ci.yml
:alt: Tests
.. image:: https://img.shields.io/pypi/v/reax.svg
:target: https://pypi.python.org/pypi/reax/
:alt: Latest Version
.. image:: https://img.shields.io/pypi/wheel/reax.svg
:target: https://pypi.python.org/pypi/reax/
.. image:: https://img.shields.io/pypi/pyversions/reax.svg
:target: https://pypi.python.org/pypi/reax/
.. image:: https://img.shields.io/pypi/l/reax.svg
:target: https://pypi.python.org/pypi/reax/
REAX — Scalable, flexible training for JAX, inspired by the simplicity of PyTorch Lightning.
REAX - Scalable Training for JAX
================================
REAX is a minimal and high-performance framework for training JAX models, designed to simplify
research workflows. Inspired by PyTorch Lightning, it brings similar high-level abstractions and
scalability to JAX users, making it easier to scale models across multiple GPUs with minimal
boilerplate. 🚀
A Port of PyTorch Lightning to JAX
----------------------------------
Much of REAX is built by porting the best practices and abstractions of **PyTorch Lightning** to
the **JAX** ecosystem. If you're familiar with PyTorch Lightning, you'll recognize concepts like:
- Training loops ⚡
- Multi-GPU training 🖥️
- Logging and checkpointing 💾
However, REAX has been designed with JAX-specific optimizations, ensuring high performance without
sacrificing flexibility.
Why REAX? 🌟
------------
- **Scalable**: Built to leverage JAX’s parallelism and scalability. ⚡
- **Minimal Boilerplate**: Simplifies the training process with just enough structure. 🧩
- **Familiar**: For users who have experience with frameworks like PyTorch Lightning, the
transition to REAX is seamless. 🔄
Installation 🛠️
---------------
To install REAX, run the following command:
.. code-block:: shell
pip install reax
REAX example
------------
Define the training workflow. Here's a toy example:
.. code-block:: python
# main.py
# ! pip install torchvision
from functools import partial
import jax, optax, reax, flax.linen as linen
import torch.utils.data as data, torchvision as tv
class Autoencoder(linen.Module):
def setup(self):
super().__init__()
self.encoder = linen.Sequential([linen.Dense(128), linen.relu, linen.Dense(3)])
self.decoder = linen.Sequential([linen.Dense(128), linen.relu, linen.Dense(28 * 28)])
def __call__(self, x):
z = self.encoder(x)
return self.decoder(z)
# --------------------------------
# Step 1: Define a REAX Module
# --------------------------------
# A ReaxModule (nn.Module subclass) defines a full *system*
# (ie: an LLM, diffusion model, autoencoder, or simple image classifier).
class ReaxAutoEncoder(reax.Module):
def __init__(self):
super().__init__()
self.ae = Autoencoder()
def setup(self, stage: "reax.Stage", batch) -> None:
if self.parameters() is None:
x = batch[0].reshape(len(batch[0]), -1)
params = self.ae.init(self.rng_key(), x)
self.set_parameters(params)
def __call__(self, *args, **kwargs):
return self.forward(*args, **kwargs)
def forward(self, x):
embedding = jax.jit(self.ae.encoder.apply)(self.parameters()["params"]["encoder"], x)
return embedding
def training_step(self, batch, batch_idx):
x = batch[0].reshape(len(batch[0]), -1)
loss, grads = jax.value_and_grad(self.loss_fn, argnums=0)(self.parameters(), x, self.ae)
self.log("train_loss", loss, on_step=True, prog_bar=True)
return loss, grads
@staticmethod
@partial(jax.jit, static_argnums=2)
def loss_fn(params, x, model):
predictions = model.apply(params, x)
return optax.losses.squared_error(predictions, x).mean()
def configure_optimizers(self):
opt = optax.adam(learning_rate=1e-3)
state = opt.init(self.parameters())
return opt, state
# -------------------
# Step 2: Define data
# -------------------
dataset = tv.datasets.MNIST(".", download=True, transform=jax.numpy.asarray)
train, val = data.random_split(dataset, [55000, 5000])
# -------------------
# Step 3: Train
# -------------------
autoencoder = ReaxAutoEncoder()
trainer = reax.Trainer(autoencoder)
trainer.fit(reax.ReaxDataLoader(train), reax.ReaxDataLoader(val))
Here, we reproduce an example from PyTorch Lightning, so we use torch vision to fetch the data,
but for real models there's no need to use this or pytorch at all.
Disclaimer ⚠️
-------------
REAX takes inspiration from PyTorch Lightning, and large portions of its core functionality are
directly ported from Lightning. If you are already familiar with Lightning, you'll feel right at
home with REAX, but we’ve tailored it to work seamlessly with JAX's performance optimizations.
Raw data
{
"_id": null,
"home_page": null,
"name": "reax",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.9",
"maintainer_email": null,
"keywords": "machine learning, jax, research",
"author": null,
"author_email": "Martin Uhrin <martin.uhrin.10@ucl.ac.uk>",
"download_url": "https://files.pythonhosted.org/packages/e1/14/3ecb1fb1b66ff230baacaa1bd217070d07ab24a76ae2936ff9f2ed893e00/reax-0.5.5.tar.gz",
"platform": null,
"description": "REAX\n====\n\n.. image:: https://codecov.io/gh/muhrin/reax/branch/develop/graph/badge.svg\n :target: https://codecov.io/gh/muhrin/reax\n :alt: Coverage\n\n.. image:: https://github.com/camml-lab/reax/actions/workflows/ci.yml/badge.svg\n :target: https://github.com/camml-lab/reax/actions/workflows/ci.yml\n :alt: Tests\n\n.. image:: https://img.shields.io/pypi/v/reax.svg\n :target: https://pypi.python.org/pypi/reax/\n :alt: Latest Version\n\n.. image:: https://img.shields.io/pypi/wheel/reax.svg\n :target: https://pypi.python.org/pypi/reax/\n\n.. image:: https://img.shields.io/pypi/pyversions/reax.svg\n :target: https://pypi.python.org/pypi/reax/\n\n.. image:: https://img.shields.io/pypi/l/reax.svg\n :target: https://pypi.python.org/pypi/reax/\n\n\nREAX \u2014 Scalable, flexible training for JAX, inspired by the simplicity of PyTorch Lightning.\n\nREAX - Scalable Training for JAX\n================================\n\nREAX is a minimal and high-performance framework for training JAX models, designed to simplify\nresearch workflows. Inspired by PyTorch Lightning, it brings similar high-level abstractions and\nscalability to JAX users, making it easier to scale models across multiple GPUs with minimal\nboilerplate. \ud83d\ude80\n\nA Port of PyTorch Lightning to JAX\n----------------------------------\n\nMuch of REAX is built by porting the best practices and abstractions of **PyTorch Lightning** to\nthe **JAX** ecosystem. If you're familiar with PyTorch Lightning, you'll recognize concepts like:\n\n- Training loops \u26a1\n- Multi-GPU training \ud83d\udda5\ufe0f\n- Logging and checkpointing \ud83d\udcbe\n\nHowever, REAX has been designed with JAX-specific optimizations, ensuring high performance without\nsacrificing flexibility.\n\nWhy REAX? \ud83c\udf1f\n------------\n\n- **Scalable**: Built to leverage JAX\u2019s parallelism and scalability. \u26a1\n- **Minimal Boilerplate**: Simplifies the training process with just enough structure. \ud83e\udde9\n- **Familiar**: For users who have experience with frameworks like PyTorch Lightning, the\n transition to REAX is seamless. \ud83d\udd04\n\nInstallation \ud83d\udee0\ufe0f\n---------------\n\nTo install REAX, run the following command:\n\n.. code-block:: shell\n\n pip install reax\n\n\nREAX example\n------------\n\nDefine the training workflow. Here's a toy example:\n\n.. code-block:: python\n\n # main.py\n # ! pip install torchvision\n from functools import partial\n import jax, optax, reax, flax.linen as linen\n import torch.utils.data as data, torchvision as tv\n\n\n class Autoencoder(linen.Module):\n def setup(self):\n super().__init__()\n self.encoder = linen.Sequential([linen.Dense(128), linen.relu, linen.Dense(3)])\n self.decoder = linen.Sequential([linen.Dense(128), linen.relu, linen.Dense(28 * 28)])\n\n def __call__(self, x):\n z = self.encoder(x)\n return self.decoder(z)\n\n\n # --------------------------------\n # Step 1: Define a REAX Module\n # --------------------------------\n # A ReaxModule (nn.Module subclass) defines a full *system*\n # (ie: an LLM, diffusion model, autoencoder, or simple image classifier).\n class ReaxAutoEncoder(reax.Module):\n def __init__(self):\n super().__init__()\n self.ae = Autoencoder()\n\n def setup(self, stage: \"reax.Stage\", batch) -> None:\n if self.parameters() is None:\n x = batch[0].reshape(len(batch[0]), -1)\n params = self.ae.init(self.rng_key(), x)\n self.set_parameters(params)\n\n def __call__(self, *args, **kwargs):\n return self.forward(*args, **kwargs)\n\n def forward(self, x):\n embedding = jax.jit(self.ae.encoder.apply)(self.parameters()[\"params\"][\"encoder\"], x)\n return embedding\n\n def training_step(self, batch, batch_idx):\n x = batch[0].reshape(len(batch[0]), -1)\n loss, grads = jax.value_and_grad(self.loss_fn, argnums=0)(self.parameters(), x, self.ae)\n self.log(\"train_loss\", loss, on_step=True, prog_bar=True)\n return loss, grads\n\n @staticmethod\n @partial(jax.jit, static_argnums=2)\n def loss_fn(params, x, model):\n predictions = model.apply(params, x)\n return optax.losses.squared_error(predictions, x).mean()\n\n def configure_optimizers(self):\n opt = optax.adam(learning_rate=1e-3)\n state = opt.init(self.parameters())\n return opt, state\n\n\n # -------------------\n # Step 2: Define data\n # -------------------\n dataset = tv.datasets.MNIST(\".\", download=True, transform=jax.numpy.asarray)\n train, val = data.random_split(dataset, [55000, 5000])\n\n # -------------------\n # Step 3: Train\n # -------------------\n autoencoder = ReaxAutoEncoder()\n trainer = reax.Trainer(autoencoder)\n trainer.fit(reax.ReaxDataLoader(train), reax.ReaxDataLoader(val))\n\nHere, we reproduce an example from PyTorch Lightning, so we use torch vision to fetch the data,\nbut for real models there's no need to use this or pytorch at all.\n\n\nDisclaimer \u26a0\ufe0f\n-------------\n\nREAX takes inspiration from PyTorch Lightning, and large portions of its core functionality are\ndirectly ported from Lightning. If you are already familiar with Lightning, you'll feel right at\nhome with REAX, but we\u2019ve tailored it to work seamlessly with JAX's performance optimizations.\n\n",
"bugtrack_url": null,
"license": null,
"summary": "REAX: A simple training framework for JAX-based projects",
"version": "0.5.5",
"project_urls": {
"Home": "https://github.com/camml-lab/reax",
"Source": "https://github.com/camml-lab/reax"
},
"split_keywords": [
"machine learning",
" jax",
" research"
],
"urls": [
{
"comment_text": null,
"digests": {
"blake2b_256": "ab29b596e9b6435a86176480c1ba970cb710d395485dba27167c4d2187640047",
"md5": "a877471545fa7bb895599baf8fbf6f7f",
"sha256": "7f5463659e2af03a65581dba29fb7f12da5eff912bbfc54464dcc396de7f7ff1"
},
"downloads": -1,
"filename": "reax-0.5.5-py3-none-any.whl",
"has_sig": false,
"md5_digest": "a877471545fa7bb895599baf8fbf6f7f",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.9",
"size": 145360,
"upload_time": "2025-07-20T19:27:48",
"upload_time_iso_8601": "2025-07-20T19:27:48.997900Z",
"url": "https://files.pythonhosted.org/packages/ab/29/b596e9b6435a86176480c1ba970cb710d395485dba27167c4d2187640047/reax-0.5.5-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": null,
"digests": {
"blake2b_256": "e1143ecb1fb1b66ff230baacaa1bd217070d07ab24a76ae2936ff9f2ed893e00",
"md5": "6e0a8caf3ef2fc2a7bb096ad9054b4a4",
"sha256": "f69dcbdf0ceb63f8d9514d33a9e74d910f9db0d324a2fa354e92d148bcbca357"
},
"downloads": -1,
"filename": "reax-0.5.5.tar.gz",
"has_sig": false,
"md5_digest": "6e0a8caf3ef2fc2a7bb096ad9054b4a4",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.9",
"size": 110661,
"upload_time": "2025-07-20T19:27:50",
"upload_time_iso_8601": "2025-07-20T19:27:50.095096Z",
"url": "https://files.pythonhosted.org/packages/e1/14/3ecb1fb1b66ff230baacaa1bd217070d07ab24a76ae2936ff9f2ed893e00/reax-0.5.5.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-07-20 19:27:50",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "camml-lab",
"github_project": "reax",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"lcname": "reax"
}
