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

[nstarman]

Adds jaxmore.nn: a scan-over-epoch training loop for JAX, plus the optional-dependency
plumbing it needs. Also demotes equinox from a required dependency to an extra.

What's new

jaxmore.nn

• AbstractScanNNTrainer — a base class for training loops built on jax.lax.scan.
  It owns the epoch loop, shuffling, batching, empty-batch skipping, and loss
  aggregation. You supply:

  • make_step and loss_agg_fn as constructor arguments;
  • init, pack_carry_state, unpack_carry_state as subclass hooks;
  • optionally prepare_data_args and prepare_step_kw, two per-epoch hooks (below).

• shuffle_and_batch — shuffles, pads to a constant batch shape, and returns a mask
  distinguishing real samples from padding. Usable samples are sorted first, so ignorable
  ones cluster at the end.

• masked_mean — mean over the True entries of a mask. Returns NaN on an empty
  mask, with a well-defined (zero) gradient.

Adapted from phasecurvefit. This part
of the library is more experimental than the rest; if you have interestingly shaped data,
please open an issue.

Per-epoch hooks

Both run once per epoch and receive epoch_idx (a traced scalar), num_epochs (a static
Python int), and epoch_key. Both default to no-ops and are skipped entirely when not
overridden.

Hook	Returns	Use it for
prepare_step_kw	a mapping, merged over step_kw	scalars: scheduled loss weights, annealed temperatures, curriculum thresholds
prepare_data_args	the data tuple, before batching	arrays with leading dim N: resampled negatives, augmentations

The split matters: anything prepare_data_args returns goes through shuffle_and_batch,
so it must be N-long. A scalar can't ride along that path — that's what prepare_step_kw
is for.

class RampTrainer(NNTrainer):
    def prepare_step_kw(self, /, *, epoch_idx, num_epochs, epoch_key):
        frac = epoch_idx / (num_epochs - 1) if num_epochs > 1 else 0.0
        return {"lam": LAM_MIN + (LAM_MAX - LAM_MIN) * frac}


def make_step(carry, batch_inputs, *, lam):  # <- arrives as a kwarg
    ...

num_epochs is passed because it is otherwise unreachable from a subclass — run() owns
it, not the instance — and a linear ramp needs the denominator.

error_if

jaxmore.error_if — a JIT/vmap-compatible runtime assertion built on jax.debug.callback,
used as the fallback when equinox isn't installed. bounded_while_loop now takes
check_termination= to turn its post-scan check off.

Dependency changes ⚠️

Package	Before	After
equinox	required	optional — jaxmore[equinox]
jax-tqdm	—	optional — jaxmore[tqdm]
jaxtyping	test-only	required (runtime)
optional_dependencies	test-only	required (runtime)

equinox is now only used to pick a better error_if in bounded_while_loop; without it,
jaxmore falls back to its own. jaxmore.nn itself has no hard Equinox dependency — it
works with any pytree carry — though the docs use Equinox + Optax because that's the common
case.

Anything currently relying on equinox arriving transitively via jaxmore will need to
depend on it directly, or install jaxmore[equinox].

Notes for reviewers

Two sharp edges are worth a look, since both are easy to get wrong and neither errors
loudly:

1. Padding is real data unless you mask it. scan needs a constant batch shape, so the
   last batch is zero-padded. A loss that ignores batch_mask trains on those fabricated
   rows — with N=100, batch_size=16 that's 12 phantom (0,0) -> 0 samples every epoch.
   The README now leads with this, and the flagship example uses masked_mean.

2. Empty batches are reachable. Because shuffle_and_batch sorts ignorable samples to
   the end, they cluster into whole batches — 40 usable of 100 at batch_size=16 yields
   four batches with no usable rows. So masked_mean has to be differentiable at an empty
   mask (it uses a clamped denominator; the naive where(count > 0, total / count, nan)
   gives NaN gradients), and run() skips such batches via lax.cond rather than burning a
   forward/backward pass on them.

Also worth a second look: the frozen-submodule example. Initializing opt_state on a
partial partition while taking gradients of the whole model hands optax a tree with slots
it never allocated. The fix — partition, differentiate the trainable half, recombine — is
non-obvious enough that it's now a documented example with a test asserting the frozen
subtree comes out bit-for-bit unchanged.

Also in this PR

• testpaths pointed at README and docs, neither of which exists (it's README.md, and
  there's no docs/), so pytest silently collected zero README tests. Fixed — every
  ```python block in the README now executes under Sybil, and several that didn't run
  have been repaired.
• New tests follow the tests/{unit,usage,benchmark} layout from 🧱 infra: set up unit, usage, and benchmark tests #18. The optimizer-strategy
  comparison moved from a time.perf_counter() loop in the unit suite to proper
  pytest-benchmark benchmarks.
• Dropped the README's "Option 1b is 5-10% faster" claim: it contradicted both the benchmark
  table further down the same page (which ranked 1b slowest) and the test's own printed
  conclusion. Measured, the three strategies sit within run-to-run noise (~1.02/1.06/1.03 ms,
  ±0.12–0.25 ms), so the docs now say to choose on clarity.
• Fixed two pre-existing PLW0108 ruff violations (in tests/unit/test_vmap.py and
  tests/benchmark/test_structured.py, both from 🧱 infra: set up unit, usage, and benchmark tests #18). They were failing ruff check and
  therefore blocking the pre-commit hook on any commit.

Known wart (not blocking)

OptDeps hits an Enum aliasing quirk: when no optional dependency is installed, every
member gets the same value (InstalledState.NOT_INSTALLED), so Python collapses them into
aliases — OptDeps.JAX_TQDM is OptDeps.EQUINOX becomes True and list(OptDeps) drops a
member. .installed still answers correctly in every install combination (verified: bare /
equinox-only / tqdm-only / both), so nothing is broken today, but .name and iteration are
unreliable in that state. Probably wants an upstream fix in optional_dependencies.

[Copilot]

Pull request overview

Adds neural-network training utilities built around jax.lax.scan, while making equinox optional and expanding the library’s internal/public API surface.

Changes:

• Introduce jaxmore.nn with AbstractScanNNTrainer, shuffle_and_batch, and masked_mean for scan-over-epoch training loops.
• Add an internal error_if helper and wire bounded_while_loop to use equinox.error_if when available (or fall back when not).
• Add a new jaxmore._src.vmap implementation and update packaging/docs for optional dependencies (equinox, jax-tqdm) and new dev/test deps (optax, etc.).

Reviewed changes

Copilot reviewed 12 out of 14 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
uv.lock	Locks new dependencies (optional deps infra, jax-tqdm, optax, etc.) and extras wiring.
pyproject.toml	Makes equinox optional via extras; adds test deps; tweaks pytest/ruff config.
src/jaxmore/__init__.py	Re-exports bounded_while_loop, vmap, error_if, and exposes nn.
src/jaxmore/nn.py	Public re-export module for NN training utilities.
src/jaxmore/_src/__init__.py	Internal re-export surface for core primitives.
src/jaxmore/_src/optional_deps.py	Central optional dependency gating via OptDeps.
src/jaxmore/_src/error.py	Adds internal error_if implemented via jax.debug.callback.
src/jaxmore/_src/while_loop.py	Adds check_termination option and conditional error-check backend selection.
src/jaxmore/_src/vmap.py	Adds enhanced vmap with static args/kwargs support and kw-axis control.
src/jaxmore/_src/nn.py	Implements AbstractScanNNTrainer, masked_mean, and shuffle_and_batch.
tests/test_unit.py	Adds tests for bounded_while_loop(check_termination=False) including jit.
tests/test_error.py	Adds unit tests for new error_if behavior under jit/vmap.
tests/test_nn_trainer.py	Adds tests for NN training utilities, including an end-to-end eqx/optax example.
README.md	Documents jaxmore.nn usage and cautions around Equinox partitioning.

Comments suppressed due to low confidence (1)

src/jaxmore/_src/while_loop.py:159

• Docstring typo: "Uses internal and error_if" reads like a grammatical error. Suggest "Uses internal error_if" (or "Uses the internal error_if").

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov]

Codecov Report

❌ Patch coverage is 98.50746% with 1 line in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@1f7733d). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
src/jaxmore/_src/while_loop.py	83.33%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main      #11   +/-   ##
=======================================
  Coverage        ?   95.55%           
=======================================
  Files           ?        9           
  Lines           ?      225           
  Branches        ?        0           
=======================================
  Hits            ?      215           
  Misses          ?       10           
  Partials        ?        0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

Copilot reviewed 15 out of 16 changed files in this pull request and generated 2 comments.