Metadata-Version: 2.1
Name: reax
Version: 0.3.0
Summary: REAX: A simple training framework for JAX-based projects
Keywords: machine learning,jax,research
Author-email: Martin Uhrin <martin.uhrin.10@ucl.ac.uk>
Requires-Python: >=3.9
Description-Content-Type: text/x-rst
Classifier: Development Status :: 4 - Beta
Classifier: License :: OSI Approved :: GNU Lesser General Public License v3 or later (LGPLv3+)
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Dist: beartype
Requires-Dist: einops
Requires-Dist: clu
Requires-Dist: flax
Requires-Dist: fsspec
Requires-Dist: jax
Requires-Dist: jaxlib
Requires-Dist: jaxtyping
Requires-Dist: lightning-utilities>=0.10.0
Requires-Dist: mlflow
Requires-Dist: optax
Requires-Dist: orbax-checkpoint
Requires-Dist: pytray
Requires-Dist: equinox
Requires-Dist: tensorboardX
Requires-Dist: tqdm
Requires-Dist: typing-extensions
Requires-Dist: black ; extra == "dev"
Requires-Dist: cloudpickle ; extra == "dev"
Requires-Dist: flit ; extra == "dev"
Requires-Dist: ipython ; extra == "dev"
Requires-Dist: pandas ; extra == "dev"
Requires-Dist: pytest ; extra == "dev"
Requires-Dist: pytest-cov ; extra == "dev"
Requires-Dist: pytest-env ; extra == "dev"
Requires-Dist: pre-commit ; extra == "dev"
Requires-Dist: pylint ; extra == "dev"
Requires-Dist: docutils ; extra == "docs"
Requires-Dist: jupyter ; extra == "docs"
Requires-Dist: nbsphinx ; extra == "docs"
Requires-Dist: pandoc ; extra == "docs"
Requires-Dist: scikit-learn ; extra == "docs"
Requires-Dist: sphinx ; extra == "docs"
Requires-Dist: sphinx-autobuild ; extra == "docs"
Requires-Dist: lightning ; extra == "examples"
Requires-Dist: pillow ; extra == "examples"
Requires-Dist: torch ; extra == "examples"
Project-URL: Home, https://github.com/muhrin/reax
Project-URL: Source, https://github.com/muhrin/reax
Provides-Extra: dev
Provides-Extra: docs
Provides-Extra: examples


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/muhrin/reax/actions/workflows/ci.yml/badge.svg
    :target: https://github.com/muhrin/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: A simple training framework for JAX-based projects

REAX is based on PyTorch Lightning and tries to bring a similar level of easy-of-use and
customizability to the world of training JAX models. Much of lightning's API has been adopted
with some modifications being made to accommodate JAX's pure function based approach.


Quick start
-----------

.. 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 LightningModule
    # --------------------------------
    # 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.
Run the model on the terminal


.. code-block:: bash

    pip install reax torchvision
    python main.py

