feat(export): ONNX for the classifier arms and the CRF encoder

.github/workflows/ci.yml

@@ -117,8 +117,12 @@ jobs:
           # don't import snakemake/plotnine). Full install (Rust extension +
           # escapepod); fall back without the Rust extension (leech_core) but keep
           # escapepod (a PyPI wheel), which leech imports unconditionally.
-          uv pip install --torch-backend=cpu -e ".[test,rust,pod5]" \
-            || uv pip install --torch-backend=cpu -e ".[test,pod5]"
+          # `onnx` is here so the export round-trip tests actually RUN rather
+          # than skipping: they compare onnxruntime against torch across the
+          # serialization boundary, which is the check an in-process assert
+          # cannot make.
+          uv pip install --torch-backend=cpu -e ".[test,rust,pod5,onnx]" \
+            || uv pip install --torch-backend=cpu -e ".[test,pod5,onnx]"
 
       - name: Run tests
         env:

CHANGELOG.md

@@ -9,6 +9,36 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- **ONNX export, for the classifier arms and the CRF encoder** (#217).
+  `leech model export --format onnx` beside the existing `--format torch`
+  (unchanged default), and `leech.crf.export.export_crf_onnx`. `torch.export`
+  makes a model loadable by anything with PyTorch and by nothing else; a runtime
+  consuming ONNX — which is what escapepod-rs runs — could not load a leech
+  model at all.
+
+  Both use the **dynamo exporter at opset 18**. `dynamo=False`, the obvious
+  first attempt, fails on these architectures with an `adaptive_avg_pool1d`
+  error that reads like a model problem and is an exporter limitation; that is
+  documented where someone will hit it, and a regression test pins it.
+
+  Each export writes a **contract** beside the graph, carrying the two things a
+  consumer needs and cannot recover from it: which input is which (including
+  that the `signal_kmer` sequence input is built in the dataset, not the model,
+  and that `leech-core` ships that encoder), and what the output means — a
+  single BCE logit, not a two-class softmax. The CRF's contract additionally
+  carries standardisation, which is in neither the config nor the checkpoint,
+  and its emitted references (`target[state_len:]`), computed from the
+  `state_len` the encoder declares.
+
+  Verified across the serialization boundary rather than in process:
+  onnxruntime against torch, 3.58e-07 for the CRF encoder against a float32 eps
+  of 1.19e-07.
+
+  New `onnx` extra (`onnx`, `onnxruntime`, `onnxscript`). CI installs it so the
+  round-trip tests run rather than skip.
+
+### Added
+
 - **`leech.crf.training`: a CTC-CRF trainer.** `CrfTrainer` runs the schedule
   and writes `model.pt` plus a `model.json` sidecar. The sidecar is not optional:
   the standardisation constants live in neither the architecture config nor the

CLAUDE.md

@@ -491,6 +491,63 @@ scans in `loss.py` are the readable reference the tests check it against, and
 the Triton kernels check against those. Keep all three — the fallback chain is
 what makes a wrong kernel visible.
 
+### ONNX export: one exporter, and what a graph cannot carry
+
+`leech.onnx_export` serves both the classifier arms (`leech model export
+--format onnx`) and the CRF encoder (`leech.crf.export`). It exists because
+`torch.export` makes a model loadable by anything with PyTorch and by nothing
+else — escapepod-rs consumes ONNX and could not load a leech model at all.
+
+**Always the dynamo exporter, opset 18.** `dynamo=False` is the obvious first
+attempt and fails on these architectures with `Unsupported: ONNX export of
+operator adaptive_avg_pool1d`. That is an exporter limitation, not a model
+problem, but the message reads like one; `tests/test_onnx_export.py` pins the
+working path so a regression to the legacy exporter says so.
+
+**Every export writes a contract beside the graph**, because two things a
+consumer needs are invisible in it: which input is which (arity and channel
+counts are visible, roles are not), and what the output means — leech
+classifiers emit a *single BCE logit*, not a two-class softmax, and reading it
+as the latter makes every call wrong without erroring. The CRF's contract adds
+standardisation (in neither the config nor the checkpoint) and the emitted
+references (`target[state_len:]`, computed from the `state_len` the encoder
+declares so no caller can pass full-length targets by hand).
+
+**The `signal_kmer` sequence input is not in the graph, and that is deliberate.**
+It is `encode_signal_kmer` output — a scatter of the one-hot k-mer context along
+the signal axis, built in the *dataset* from the base-to-signal map. Two options
+were on the table: bake it into the graph as an ONNX prefix, or keep it a
+documented call. **It stays a documented call**, because the scatter needs the
+base-to-signal map, which comes from the move table — an ONNX prefix would still
+have to take that map as an input, so it moves where the scatter happens without
+removing the consumer's obligation, and it would bake `signal_kmer_context` into
+the graph. The contract names the function and its parameters rather than
+leaving them to be re-derived — the same choice escapepod-rs's charging bundle
+makes by carrying its recipe in `metadata.json`.
+
+**But "one definition" is not yet true, and the fix is upstream.** `leech-core`
+ships the encoder, and it is tempting to say a consumer should just call it —
+except `leech_core` is `crate-type = ["cdylib"]`, a Python extension module, so
+escapepod-rs cannot link it. The primitive itself
+(`rust/src/encoding.rs::encode_signal_kmer_inner`) is pure, dependency-free and
+carries no model vocabulary: sequence ints, a base-to-signal map, a signal
+length and a k-mer context in, a `(4 * kmer_len, signal_len)` scatter out. That
+is an `escapepod-signal` primitive by every rule this stack already applies —
+and escapepod-signal owns `mapping`, which *produces* the map this consumes, so
+today the producer is upstream and the consumer is downstream. Moving it is
+rnabioco/escapepod-rs#271. Until it does,
+any Rust consumer has to transcribe it, which is a second definition, and the
+mitigation is a cross-language golden of the kind the CRF decode and the
+charging features already have.
+
+**Verification crosses the serialization boundary.** `verify_onnx` runs
+onnxruntime against torch and returns the max absolute difference, which is what
+an in-process assert cannot check. The CRF encoder measures 3.58e-07 against a
+float32 eps of 1.19e-07; the production classifier arms measured 4.77e-07 and
+1.19e-06 (rnabioco/leech#217). It also pins onnxruntime's thread pool, whose
+default affinity call fails inside any cgroup-restricted allocation and floods
+stderr.
+
 ### Key Classes and Functions
 
 **`MoveTable` (features.py)**
@@ -773,9 +830,11 @@ The codebase is feature-complete (v0.7.0):
 - ✓ CTC-CRF sequence models (`leech.crf`): encoder, training objective with an
   analytic forward-backward, optional Triton lattice kernels, and the two-pass
   Viterbi decode — ported from escapepod-models, torch + numpy only. Plus the
-  manifest seam, the corpus builder (`plan_corpus`/`build_corpus`) and the
-  trainer (`CrfTrainer`). The ONNX export and the metrics/eval half are not
-  here yet.
+  manifest seam, the corpus builder (`plan_corpus`/`build_corpus`), the trainer
+  (`CrfTrainer`) and the ONNX export. The metrics/eval half is not here yet.
+- ✓ ONNX export for the classifier arms and the CRF encoder
+  (`leech model export --format onnx`, `leech.crf.export`), dynamo exporter at
+  opset 18, each with a contract sidecar and a round-trip check against torch
 
 All core functionality is implemented and ready for use.
 

docs/api/crf.md

@@ -162,6 +162,36 @@ AUROC/F1 checkpointing. Five things it does that are easy to get wrong:
   gradient counts. An epoch mean alone cannot tell one blown batch from a
   thousand mediocre ones.
 
+## Exporting for a native runtime
+
+`export_crf_onnx` writes `crf_encoder.onnx` plus a `metadata.json` contract.
+The **encoder only** — the decode is not expressible in standard ONNX ops, which
+is why `escapepod-demux` owns it.
+
+```python
+from leech.crf.export import export_crf_onnx
+
+export_crf_onnx("run/", "export/", sidecar="run/", references={"code01": target})
+```
+
+```
+input   signal  [batch, 1, chunk]                 float32, BATCH-major
+output  scores  [chunk // stride, batch, n_score] float32, TIME-major
+```
+
+Time-major output is the trap: the boundary CNN in the same stack is batch-major
+`[B, 2, L]`, so a consumer reusing that assumption silently transposes rather
+than failing, and needs its own load-time shape probe.
+
+The sidecar is not decoration. **Standardisation is in neither the architecture
+config nor the checkpoint** — the trainer derives it from the corpus — so a
+consumer holding only weights cannot standardise and decodes silently worse.
+Passing `references=` writes what the model *emits* (`target[state_len:]`),
+computed once from the `state_len` the encoder declares, so no caller can supply
+full-length targets by hand and inflate every edit distance.
+
+Requires the `onnx` extra: `uv sync --extra onnx`.
+
 ## Encoder
 
 ::: leech.crf.encoder.CrfEncoder
@@ -269,6 +299,16 @@ AUROC/F1 checkpointing. Five things it does that are easy to get wrong:
     options:
       show_root_heading: true
 
+## Export API
+
+::: leech.crf.export.export_crf_onnx
+    options:
+      show_root_heading: true
+
+::: leech.crf.export.load_training_sidecar
+    options:
+      show_root_heading: true
+
 ## Configuration
 
 ::: leech.crf.config.load_config

pyproject.toml

@@ -63,6 +63,16 @@ rust = ["leech-core==0.8.0"]
 # Retained as a no-op alias: escapepod is now a required dependency, but
 # `--extra pod5` appears in existing docs, CI, and user scripts.
 pod5 = []
+# ONNX export (`leech model export --format onnx`, `leech.crf.export`).
+# Optional because torch alone can export the graph -- it is the round-trip
+# VERIFY, which compares onnxruntime against torch across the serialization
+# boundary, that needs these. That check is what catches what an in-process
+# assert cannot, so the extra is worth having rather than skipping the step.
+onnx = [
+  "onnx>=1.17.0",
+  "onnxruntime>=1.20.0",
+  "onnxscript>=0.1.0",  # required by torch.onnx.export(dynamo=True)
+]
 # Minimal deps to run the test suite (used by CI's test job).
 test = [
   "pytest>=8.0.0",
@@ -72,6 +82,7 @@ test = [
 # Full local dev env: test runners + lint/type/QA tooling.
 dev = [
   "leech[test]",
+  "leech[onnx]",
   "ruff>=0.16.1",  # Modern linter/formatter (replaces black + flake8)
   "ty>=0.0.66",
   "ruff>=0.16.1",  # Modern linter/formatter (replaces black + flake8)

src/leech/cli.py

@@ -1062,17 +1062,35 @@ def calibrate(model_dir, val_data, device, batch_size, num_workers, method, reg_
     "-o",
     required=True,
     type=click.Path(path_type=Path),
-    help="Output TorchScript .pt file path",
-)
-def export(model_dir, output):
-    """Export a trained model as a standalone .pt file.
-
-    The exported file is loadable with just torch.export.load() — no leech
-    codebase required. Model config is embedded in the file.
+    help="Output file path (.pt for torch, .onnx for onnx)",
+)
+@click.option(
+    "--format",
+    "fmt",
+    type=click.Choice(["torch", "onnx"]),
+    default="torch",
+    show_default=True,
+    help="torch: loadable with torch.export.load(). onnx: loadable by any ONNX "
+    "runtime, and the only option for a consumer without PyTorch.",
+)
+def export(model_dir, output, fmt):
+    """Export a trained model as a standalone file.
+
+    \b
+    torch  loadable with torch.export.load() — no leech required; config is
+           embedded in the file.
+    onnx   loadable by any ONNX runtime. Writes a `.json` contract beside the
+           graph naming each input's role and the output convention (a single
+           BCE logit, NOT a two-class softmax), neither of which a consumer can
+           recover from the graph.
+
+    With `--seq-encoding signal_kmer`, the `sequence` input is produced outside
+    the model, in the dataset. The contract names how; `leech-core` ships that
+    encoder, and calling it keeps one definition of the rule.
     """
     from leech.commands.bundle import handle_export
 
-    handle_export(model_dir=model_dir, output=output)
+    handle_export(model_dir=model_dir, output=output, fmt=fmt)
 
 
 @model.command()

src/leech/commands/bundle.py

@@ -237,7 +237,7 @@ def handle_bundle_info(bundle: Path) -> dict[str, Any]:
     return metadata
 
 
-def handle_export(model_dir: Path, output: Path) -> Path:
+def handle_export(model_dir: Path, output: Path, fmt: str = "torch") -> Path:
     """
     Handle the export command logic.
 
@@ -250,7 +250,12 @@ def handle_export(model_dir: Path, output: Path) -> Path:
     """
     from leech.model_export import export_single_model
 
-    output_path = export_single_model(model_dir, output)
+    if fmt == "onnx":
+        from leech.model_export import export_single_model_onnx
+
+        output_path = export_single_model_onnx(model_dir, output)
+    else:
+        output_path = export_single_model(model_dir, output)
     size_mb = output_path.stat().st_size / (1024 * 1024)
 
     table = Table(title="Export Summary", show_header=True, header_style="bold magenta")