✨ feat(nn): AbstractScanNNTrainer for scan-over-epoch NN training

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

Author: nstarman

README.md

@@ -148,3 +148,171 @@ result = bounded_while_loop(
 )
 print(result)  # (Array(3, dtype=int32), Array(8, dtype=int32))
 ```
+
+### `jaxmore.nn` — efficient neural network training with JAX scan
+
+The `AbstractScanNNTrainer` class provides a foundation for building efficient
+training loops using
+[`jax.lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html)
+to scan over batches and epochs. It handles shuffling, batching, and loss
+aggregation automatically, while you focus on defining `make_step()`.
+
+> **Note**: The neural network training utilities in `jaxmore.nn` are more
+> experimental than other components of this library. The code is adapted from
+> [phasecurvefit](https://phasecurvefit.readthedocs.io/en/latest/). If you have
+> interestingly shaped data or encounter errors, please
+> [submit an issue](https://github.com/GalacticDynamics/jaxmore/issues) so we
+> can continue to generalize this training code.
+
+Here's a complete example of training a simple feed-forward network:
+
+```python
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+import optax
+from jaxmore.nn import AbstractScanNNTrainer, masked_mean
+
+
+# Define a simple neural network using eqx.nn.MLP
+model = eqx.nn.MLP(
+    in_size=2,
+    out_size=1,
+    width_size=32,
+    depth=2,
+    activation=jax.nn.relu,
+    key=jr.key(0),
+)
+
+
+# Create a trainer subclass
+class NNTrainer(AbstractScanNNTrainer):
+    """Concrete trainer implementation for SimpleNN."""
+
+    def init(self, *, key, X, y, learning_rate=1e-2):
+        """Initialize model and training data."""
+        model_key, data_key = jr.split(key)
+        model = eqx.nn.MLP(
+            in_size=X.shape[1],
+            out_size=1,
+            width_size=32,
+            depth=2,
+            activation=jax.nn.relu,
+            key=model_key,
+        )
+
+        optimizer = optax.adam(learning_rate)
+        opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
+
+        # Create training carry: (model, optimizer_state, rng_key)
+        carry_key = jr.fold_in(data_key, 0)
+        initial_carry = (model, opt_state, carry_key)
+
+        # All samples are usable (True), none are padding (False)
+        mask = jnp.ones(len(X), dtype=bool)
+        epoch_data = (mask, (X, y))
+
+        return initial_carry, epoch_data
+
+    def pack_carry_state(self, carry):
+        """Partition model into arrays and static structure."""
+        model, opt_state, key = carry
+        model_dyn, model_static = eqx.partition(model, eqx.is_array)
+        return (model_dyn, opt_state, key), {"model_static": model_static}
+
+    def unpack_carry_state(self, carry, static):
+        """Reconstruct full model from partitioned state."""
+        model_dyn, opt_state, key = carry
+        model_static = static["model_static"]
+        model = eqx.combine(model_dyn, model_static)
+        return (model, opt_state, key)
+
+
+# Define optimizer
+optimizer = optax.adam(1e-2)
+
+
+# Define the per-batch training step
+def make_step(carry, batch_inputs):
+    """Execute one batch of training."""
+    model, opt_state, key = carry
+    batch_mask, (X_batch, y_batch) = batch_inputs
+
+    def loss_fn(model):
+        preds = jax.vmap(model)(X_batch)
+        mse = jnp.mean((preds.squeeze() - y_batch) ** 2)
+        return mse
+
+    loss, grads = eqx.filter_value_and_grad(loss_fn)(model)
+    updates, opt_state = optimizer.update(grads, opt_state)
+    model = eqx.apply_updates(model, updates)
+
+    return loss, (model, opt_state, key)
+
+
+# Generate synthetic training data
+key = jr.key(0)
+X = jr.normal(key, (100, 2))
+y = 2 * X[:, 0] + X[:, 1] + 0.1 * jr.normal(jr.fold_in(key, 1), (100,))
+
+# Create trainer and initialize
+trainer = NNTrainer(make_step=make_step, loss_agg_fn=masked_mean)
+carry, epoch_data = trainer.init(key=key, X=X, y=y, learning_rate=1e-2)
+
+# Train for 10 epochs with batch size 16
+final_carry, losses = trainer.run(
+    carry,
+    epoch_data,
+    num_epochs=10,
+    batch_size=16,
+    key=key,
+    show_pbar=False,
+)
+
+print(f"Final epoch loss: {losses[-1]:.6f}")  # doctest: +SKIP
+```
+
+Key features:
+
+- **Automatic batching** — `shuffle_and_batch()` handles shuffling and padding
+- **Efficient scanning** — Uses `jax.lax.scan` for epochs and batches
+  (JAX-friendly)
+- **Model partitioning** — Separate dynamic arrays (model weights) from static
+  structure for efficient JIT compilation
+- **Loss aggregation** — Customize how per-batch losses combine into epoch
+  losses
+
+#### ⚠️ Sharp Edge: Equinox Models Require `eqx.partition`
+
+**CRITICAL**: When using Equinox models (`eqx.nn.MLP`, `eqx.nn.Linear`, etc.),
+you **MUST** implement `pack_carry_state()` and `unpack_carry_state()` using
+`eqx.partition` and `eqx.combine`. This is not optional!
+
+**Why?** Equinox modules contain methods decorated with `@jax.custom_jvp`, which
+JAX cannot scan over directly. Without partitioning, you'll encounter:
+
+```
+TypeError: Argument '...' is not a valid JAX type
+```
+
+**Correct pattern** (as shown above):
+
+```python
+def pack_carry_state(self, carry):
+    model, opt_state, key = carry
+    # Separate arrays from static structure (methods, activations, etc.)
+    model_dyn, model_static = eqx.partition(model, eqx.is_array)
+    return (model_dyn, opt_state, key), {"model_static": model_static}
+
+
+def unpack_carry_state(self, carry, static):
+    model_dyn, opt_state, key = carry
+    model_static = static["model_static"]
+    # Reconstruct full model
+    model = eqx.combine(model_dyn, model_static)
+    return (model, opt_state, key)
+```
+
+This separates trainable arrays (which JAX can scan over) from static Python
+objects (which it cannot), enabling the training loop to work correctly.

pyproject.toml

@@ -19,7 +19,7 @@ classifiers = [
   "Topic :: Scientific/Engineering",
   "Typing :: Typed",
 ]
-dependencies = ["equinox>=0.11.0", "jax>=0.4.20"]
+dependencies = ["jax>=0.4.20", "jaxtyping>=0.3.7", "optional_dependencies>=0.3"]
 description = "There's more to JAX."
 dynamic = ["version"]
 license.file = "LICENSE"
@@ -33,6 +33,10 @@ Changelog     = "https://github.com/GalacticDynamics/jaxmore/releases"
 Discussions   = "https://github.com/GalacticDynamics/jaxmore/discussions"
 Homepage      = "https://github.com/GalacticDynamics/jaxmore"
 
+[project.optional-dependencies]
+equinox = ["equinox>=0.11.0"]
+tqdm    = ["jax-tqdm>=0.1.0"]
+
 
 [build-system]
 build-backend = "hatchling.build"
@@ -61,12 +65,14 @@ lint = ["mypy>=1.19.0", "pre-commit>=4.1.0", "pylint>=3.3.8"]
 nox = ["nox>=2024.10.9", "nox-uv>=0.6.3"]
 test = [
   "beartype>=0.19.0",
+  "equinox>=0.11.0",
+  "jax-tqdm>=0.1.0",
   "jaxtyping>=0.2.34",
-  "optional_dependencies>=0.3",
+  "optax>=0.2.0",
   "pytest>=8.4.2",
   "pytest-cov>=6.2.1",
   "pytest-github-actions-annotate-failures>=0.3.0",
-  "sybil[pytest]>=9.2.0",
+  "sybil>=9.2.0",
 ]
 
 
@@ -97,7 +103,10 @@ addopts = [
   "-p no:doctest",                                                                                                     # using sybil
   "-ra",
 ]
-filterwarnings = ["error"]
+filterwarnings = [
+  "error",
+  "ignore:Setting `jax_pmap_shmap_merge` is deprecated:DeprecationWarning",
+]
 log_cli_level = "INFO"
 minversion = "8.3"
 testpaths = ["src/", "tests/", "README", "docs"]
@@ -140,11 +149,22 @@ ignore = [
   "ANN401",  # Dynamically typed expressions (typing.Any) are disallowed
   "TD002",   # Missing author in TODO
   "TD003",   # Missing issue link on the line following this TODO
+  "PYI041",  # Use `X | Y` instead of `Union[X, Y]` in function signatures
 ]
 
 [tool.ruff.lint.per-file-ignores]
 "noxfile.py" = ["T20"]
-"tests/**"   = ["ANN", "E731", "INP001", "S101", "T20"]
+"tests/**" = [
+  "ANN",
+  "E731",
+  "FBT003",
+  "INP001",
+  "N803",
+  "N806",
+  "RET504",
+  "S101",
+  "T20",
+]
 
 [tool.ruff.lint.isort]
 combine-as-imports     = true
@@ -158,6 +178,10 @@ ignore-paths = [".*/_version.py"]
 messages_control.disable = [
   "design",
   "fixme",
+  "import-error",
+  "import-outside-toplevel",
+  "invalid-enum-extension",
+  "invalid-name",
   "line-too-long",
   "missing-function-docstring",
   "missing-module-docstring",

src/jaxmore/__init__.py

@@ -1,6 +1,6 @@
 """Copyright (c) 2026 Nathaniel Starkman. All rights reserved."""
 
-__all__ = ("bounded_while_loop", "vmap")
+__all__ = ("bounded_while_loop", "nn", "vmap")
 
-from ._vmap import vmap
-from ._while import bounded_while_loop
+from . import nn
+from ._src import bounded_while_loop, vmap

src/jaxmore/_src/__init__.py

@@ -0,0 +1,5 @@
+"""Copyright (c) 2026 Nathaniel Starkman. All rights reserved."""
+
+from .error import *  # noqa: F403
+from .vmap import *  # noqa: F403
+from .while_loop import *  # noqa: F403