Cleanup

.github/workflows/ci.yaml

@@ -20,7 +20,7 @@ jobs:
       - precommit
     strategy:
       matrix:
-        python-version: [ 3.12 ]
+        python-version: [ 3.13 ]
     steps:
       - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}
@@ -32,18 +32,18 @@ jobs:
           version: "latest"
       - name: Install dependencies
         run: |
-          uv sync --dev
+          uv sync --all-groups --all-extras
       - name: Run lints
         run: |
-          make lints
+          uv run ruff check blaxbird examples
 
   tests:
     runs-on: ubuntu-latest
     needs:
       - lints
     strategy:
       matrix:
-        python-version: [ 3.12 ]
+        python-version: [ 3.13 ]
     steps:
       - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}
@@ -55,11 +55,11 @@ jobs:
           version: "latest"
       - name: Install dependencies
         run: |
-          uv sync --dev
+          uv sync --all-extras
       - name: Run tests
         run: |
-          make tests
+          uv run pytest
       - name: Upload coverage reports to Codecov
-        uses: codecov/codecov-action@v5
+        uses: codecov/codecov-action@v3
         env:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

.github/workflows/release.yaml

@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.11]
+        python-version: [ 3.13 ]
     steps:
     - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}

.gitignore

@@ -155,3 +155,6 @@ cython_debug/
 .vscode/
 
 .DS_Store
+
+# git worktrees for isolated feature work
+.worktrees/

.gitlint

@@ -1,8 +1,15 @@
 [general]
 ignore=body-is-missing
+contrib=contrib-title-conventional-commits
 
 [title-min-length]
 min-length=10
 
+[title-max-length]
+line-length=72
+
+[body-max-line-length]
+line-length=72
+
 [title-must-not-contain-word]
 words=wip,todo

.pre-commit-config.yaml

@@ -18,7 +18,7 @@ repos:
   hooks:
   - id: mypy
     args: ["--ignore-missing-imports"]
-    files: "(blaxbird)"
+    files: "(blaxbird|examples)"
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
   rev: v0.3.0

.python-version

@@ -1 +1 @@
-3.11.9
+3.12.10

README.md

@@ -5,23 +5,15 @@
 
 > A high-level API to build and train NNX models
 
-## About
-
-`Blaxbird` [blækbɜːd] is a high-level API to easily build NNX models and train them on CPU or GPU.
+`Blaxbird` [blæksbɜːd] is a high-level API to easily build NNX models and train them on CPU or GPU.
 
 Using `blaxbird` one can
 - concisely define models and loss functions without the usual JAX/Flax verbosity,
 - easily define checkpointers that save the best and most current network weights,
 - distribute data and model weights over multiple processes or GPUs,
 - define hooks that are periodically called during training.
 
-In addition, `blaxbird` offers high-quality implementations of common neural network modules and algorithms, such as:
-
-- MLPs, DiTs, UNets,
-- Flow Matching and Denoising Score Matching (EDM schedules) models with Euler and Heun samplers,
-- Consistency Distillation/Matching models.
-
-## Example
+## Quickstart
 
 To use `blaxbird`, one only needs to define a model, a loss function, and train and validation step functions:
 ```python
@@ -62,20 +54,37 @@ train = train_fn(
   eval_every_n_steps=10,
   n_eval_batches=10
 )
-train(jr.key(2), model, optimizer, train_itr, val_itr)
+train(jr.key(2), optimizer, train_itr, val_itr)
 ```
 
-See the entire self-contained example in [examples/mnist_classification](examples/mnist_classification).
+## Examples
+
+Full self-contained examples (flow matching, small LMs training, ...), can be found in [examples](examples/).
 
-## Usage
+## Installation
+
+To install the package from PyPI, call:
+
+```bash
+pip install blaxbird
+```
+
+To install the latest GitHub <RELEASE>, just call the following on the command line:
+
+```bash
+pip install git+https://github.com/dirmeier/blaxbird@<RELEASE>
+```
+
+## API
 
 `train_fn` is a higher order function with the following signature:
 
 ```python
 def train_fn(
   *,
   fns: tuple[Callable, Callable],
-  shardings: Optional[tuple[jax.NamedSharding, jax.NamedSharding]] = None,
+  mesh: jax.sharding.Mesh | None = None,
+  data_partition_spec: jax.sharding.PartitionSpec = jax.sharding.PartitionSpec(),
   n_steps: int,
   eval_every_n_steps: int,
   n_eval_batches: int,
@@ -85,6 +94,8 @@ def train_fn(
   ...
 ```
 
+The returned `train` callable has signature `train(rng_key, optimizer, train_itr, val_itr) -> None` -- it derives the model from `optimizer.model`, so there is no separate `model` argument.
+
 We briefly explain the more ambiguous argument types below.
 
 ### `fns`
@@ -109,26 +120,33 @@ The loss function that is called by both computes a *scalar* loss value. B
 While `train_step` returns has to return the loss and gradients, `val_step` only needs
 to return the loss.
 
-### `shardings`
+### `mesh` and `data_partition_spec`
 
 To specify how data and model weights are distributed over devices and processes,
 `blaxbird` uses JAX' [sharding](https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) functionality.
 
-`shardings` is again specified by a tuple, one for the model sharding, the other for the data sharding.
-An example is shown below, where we only distributed the data over `num_devices` devices.
-You can, if you don't want to distribute anything, just set the argument to `None` or not specify it.
+`mesh` is a `jax.sharding.Mesh` describing your device topology. Per-parameter sharding is derived from each parameter's own `nnx.with_partitioning`
+annotation (see the [flax.nnx docs](https://flax.readthedocs.io/en/latest/)) via `nnx.get_named_sharding` -- parameters without such an annotation default to fully replicated, so a mesh with no annotated parameters gives plain data parallelism. `data_partition_spec` controls how each training/eval batch is sharded across `mesh` (defaults to `PartitionSpec()`, fully replicated). You can, if you don't want to distribute anything, just leave `mesh` as `None` or not specify it.
+
+An example is shown below, sharding only the data over `num_devices` devices
+(the model has no `with_partitioning` annotations, so it stays fully
+replicated):
 
 ```python
-def get_sharding():
+def get_mesh():
   num_devices = jax.local_device_count()
-  mesh = jax.sharding.Mesh(
+  return jax.sharding.Mesh(
     mesh_utils.create_device_mesh((num_devices,)), ("data",)
   )
-  model_sharding = jax.NamedSharding(mesh, jax.sharding.PartitionSpec())
-  data_sharding = jax.NamedSharding(mesh, jax.sharding.PartitionSpec("data"))
-  return model_sharding, data_sharding
+
+mesh = get_mesh()
 ```
 
+Pass `mesh=mesh, data_partition_spec=jax.sharding.PartitionSpec("data")` to
+`train_fn`. For real FSDP/tensor-parallel sharding, annotate your model's
+layers with `nnx.with_partitioning` -- see
+[examples/fsdp_tp_demo](examples/fsdp_tp_demo) for a worked 2D-mesh example.
+
 ### `hooks`
 
 `hooks` is a list of callables which are periodically called during training.
@@ -205,6 +223,53 @@ hook_save, *_ = get_default_checkpointer(
 )
 ```
 
+#### An EMA `hook`
+
+We provide a hook for tracking an exponential moving average (EMA) of a
+model's weights, constructed via `get_ema_hook`.
+
+The signature is:
+
+```python
+def get_ema_hook(
+  model: nnx.Module, decay: float = 0.999
+) -> tuple[Callable, Callable]
+```
+
+Its arguments are:
+- `model`: the model from which the EMA state is initialized.
+- `decay`: the EMA decay rate.
+
+It returns a tuple `(hook_fn, get_ema_model_fn)`:
+- `hook_fn(step, *, model, **kwargs) -> None`: updates the tracked EMA
+  weights every training step.
+- `get_ema_model_fn(model: nnx.Module) -> nnx.Module`: returns a new,
+  independent `nnx.Module` with the same structure as `model` but with the
+  tracked EMA parameter values.
+
+For instance, you would construct and use the EMA hook like this:
+
+```python
+from blaxbird import get_ema_hook
+
+ema_hook, get_ema_model = get_ema_hook(model, decay=0.999)
+
+train = train_fn(
+  fns=(train_step, val_step),
+  n_steps=n_steps,
+  eval_every_n_steps=eval_every_n_steps,
+  n_eval_batches=n_eval_batches,
+  hooks=[ema_hook],
+)
+train(jr.key(1), optimizer, train_itr, val_itr)
+
+ema_model = get_ema_model(optimizer.model)
+```
+
+Note: EMA state is not integrated with `get_default_checkpointer` --
+saving and restoring EMA state alongside model checkpoints is not
+covered here.
+
 ### Restoring a run
 
 You can also use `get_default_checkpointer` to restart the run where you left off.
@@ -225,44 +290,62 @@ You can then do either of:
 model = CNN(rngs=nnx.rnglib.Rngs(jr.key(1)))
 optimizer = nnx.Optimizer(model, optax.adam(1e-4))
 
-model, optimizer = restore_best(model, optimizer)
-model, optimizer = restore_last(model, optimizer)
+optimizer = restore_best(optimizer)
+optimizer = restore_last(optimizer)
 ```
 
+`restore_best`/`restore_last` take and return `optimizer` only -- the wrapped
+model (`optimizer.model`) and `opt_state` are both updated in place on the
+same optimizer instance, since `nnx.Optimizer` already owns the model it
+wraps.
+
 ### Doing training
 
-After having defined train functions, hooks and shardings, you can train your model like this:
+After having defined train functions, hooks and a mesh, you can train your model like this:
 
 ```python
 train = train_fn(
   fns=(train_step, val_step),
   n_steps=n_steps,
   eval_every_n_steps=eval_every_n_steps,
   n_eval_batches=n_eval_batches,
-  shardings=(model_sharding, data_sharding),
+  mesh=mesh,
+  data_partition_spec=jax.sharding.PartitionSpec("data"),
   hooks=hooks,
   log_to_wandb=False,
 )
-train(jr.key(1), model, optimizer, train_itr, val_itr)
+train(jr.key(1), optimizer, train_itr, val_itr)
 ```
 
-Self-contained examples that also explain how the data loaders should look like can be found
-in [examples](examples).
+## Contributing
 
-## Installation
+Contributions in the form of pull requests are more than welcome. A good way to
+start is to check out issues labelled
+[good first issue](https://github.com/dirmeier/surjectors/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22).
 
-To install the package from PyPI, call:
+In order to contribute:
 
-```bash
-pip install blaxbird
-```
+1) Clone `surjectors` and install `uv` from [here](https://docs.astral.sh/uv/getting-started/installation/).
+2) Install all dependencies using `uv sync --all-groups`.
+3) Install the Git hooks:
 
-To install the latest GitHub <RELEASE>, just call the following on the command line:
+   ```bash
+   uv run pre-commit install -t pre-commit -t commit-msg
+   ```
+4) Create a new branch locally, e.g. `git checkout -b feature/my-new-feature`.
+5) Implement your contribution and ideally a test case.
+6) Check your work (see below).
+7) Submit a PR 🙂.
 
-```bash
-pip install git+https://github.com/dirmeier/blaxbird@<RELEASE>
-```
+### Development commands
 
-## Author
+The project uses `uv` for everything:
 
-Simon Dirmeier <a href="mailto:simd@mailbox.org">simd@mailbox.org</a>
+```bash
+uv sync --all-groups
+uv run pytest
+uv run ruff format blaxbird examples
+uv run ruff check --fix blaxbird examples
+uv run mypy blaxbird examples
+uv run pre-commit run --all-files
+```

blaxbird/__init__.py

@@ -1,8 +1,9 @@
 """blaxbird: A high-level API for building and training Flax NNX models."""
 
-__version__ = "0.1.1"
+__version__ = "0.2.0"
 
 from blaxbird._src.checkpointer import get_default_checkpointer
+from blaxbird._src.hooks import get_ema_hook
 from blaxbird._src.trainer import train_fn
 
-__all__ = ["get_default_checkpointer", "train_fn"]
+__all__ = ["get_default_checkpointer", "get_ema_hook", "train_fn"]