feat(export): ONNX for the classifier arms and the CRF encoder - #220
Merged
Conversation
Closes #217. `torch.export` makes a trained model loadable by anything with PyTorch and by nothing else, so a runtime consuming ONNX — which is what escapepod-rs runs — could not load a leech model at all. That gap has been recorded downstream as though "not a shippable format" were a property of the models. It is not: the graphs convert and agree with torch to within float32 rounding. `leech model export --format onnx` beside the existing `--format torch` (unchanged default), and `leech.crf.export.export_crf_onnx`, both over a shared `leech.onnx_export`. The shared layer is justified by having two consumers, not by anticipation: the CRF encoder and the classifier arms need the same exporter settings, the same round-trip check, and the same idea of a contract. 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, output size that are not factor of input size That is an exporter limitation and not a model problem, but the message reads like one and sends whoever hits it looking in the wrong place, so it is documented where they will be and pinned by a test that exports exactly that shape. 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 `seq_channels = sum(signal_kmer_context) * 4 + 4` is not something a consumer should re-derive. And what the output means: leech classifiers emit a SINGLE BCE logit, not a two-class softmax, and read as the latter every call is wrong while nothing errors. The CRF's contract adds standardisation — in neither the architecture config nor the checkpoint, since the trainer derives it from the corpus — and its emitted references (`target[state_len:]`), computed from the `state_len` the encoder declares so no caller can hand it full-length targets and inflate every edit distance. THE SIGNAL_KMER INPUT STAYS OUTSIDE THE GRAPH, deliberately. #217 raises this as the consequential choice, so: 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. An ONNX prefix would still have to take that map as an input — the map comes from the move table — so it moves where the scatter happens without removing the consumer's obligation, and it would bake `signal_kmer_context` into the graph. leech-core already ships the encoder, so depending on it keeps one definition; a reimplementation creates a second that diverges silently, which is how a downstream repo reproduced a superseded feature definition for two months. The contract names the function and its parameters instead — the same choice escapepod-rs's charging bundle makes by carrying its recipe in metadata.json rather than in flags. Verification crosses the serialization boundary rather than asserting in process: onnxruntime against torch, 3.58e-07 for the CRF encoder against a float32 eps of 1.19e-07. It also pins onnxruntime's thread pool, whose default affinity call fails inside any cgroup-restricted allocation and floods stderr without affecting results. New `onnx` extra (onnx, onnxruntime, onnxscript — the last is what `torch.onnx.export(dynamo=True)` needs). CI installs it so the round-trip tests RUN rather than skip; a skipped round trip is exactly the check being skipped. 18 tests. Full suite 1460 passed, 44 skipped.
The signal_kmer note said leech-core ships the encoder "so depend on it rather than reimplementing". That is right for a Python consumer and wrong for the one that actually matters here: `leech_core` is `crate-type = ["cdylib"]`, a Python extension module, so escapepod-rs cannot link it. As written the note reads as though the one-definition problem were solved. It is not. State it accurately, and name the real fix. The primitive itself (`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 already owns `mapping`, which PRODUCES the map this consumes: the producer is upstream and the consumer is downstream, which is backwards. Until it moves, a Rust consumer has to transcribe the rule, and the mitigation is a cross-language golden of the kind the CRF decode and the charging features already have. The choice to keep the encoding out of the graph is unchanged and still stands on its own two reasons; only the "already solved" framing goes.
This was referenced Aug 25, 2026
jayhesselberth
added a commit
to rnabioco/escapepod-rs
that referenced
this pull request
Aug 25, 2026
Closes #271. `mapping` (#262) already **produces** a base→signal map. The primitive that **consumes** one — scattering the one-hot k-mer context along the signal axis, i.e. the 36-channel `sequence` input of a leech `seq_encoding="signal_kmer"` model — lived downstream in `leech_core`, a `cdylib` Python extension module Rust cannot link. Since that tensor is computed in the dataset it is not in leech's exported ONNX graph (rnabioco/leech#220), so a Rust runtime has to build it before it can call the model at all. Transcribing was the only option, and this stack's track record on transcribing is in the issue. ## What lands `escapepod_signal::seq_encoding`: | | | |---|---| | `encode_signal_kmer` / `_into` | the encoding itself, row-major `(4 * kmer_len, signal_len)`; the `_into` form so a loop over chunks allocates once rather than per chunk | | `sequence_ints_with_context` | cuts the `seq_len + before + after` window a chunk needs, padding off either end of the read with `UNKNOWN_BASE` — leech's `extractor.py` rule, written once | | `base_to_int` / `sequence_to_int` / `UNKNOWN_BASE` | the `A/C/G/T=U` alphabet both take. `resquiggle::kmer_table` now shares this definition instead of carrying its own copy | | `KmerContext` | names the `(before, after)` pair and owns `kmer_len()` / `channels()` | `KmerContext` is a named struct for the reason `CigarOp` is: transposing an asymmetric context displaces every k-mer window by `before - after` bases and still returns a correctly shaped tensor. Writing the tests turned up the honest version of that claim — given an already-cut `seq_ints`, the encoder only depends on the *sum*; the split matters where the context is **cut**. So the cut is now a function here too, and both the docs and `before_and_after_are_not_interchangeable` say which step is which. ## Parity `tests/signal_kmer_parity.rs` pins 35 cases bit-exactly against leech's NumPy reference (`tests/fixtures/gen_signal_kmer_golden.py`, same generator pattern as `gen_kmer_levels_golden.py`; numpy-only, no `leech_core` build). The encoding is exactly zeros and ones, so there is no tolerance to argue about. Cases cover the default `(4, 4)` context, asymmetric contexts, unknown bases, uracil, empty spans, spans off both ends, a zero-width window, and a seeded random sweep. The golden comes from the **NumPy** path deliberately, because leech's two implementations of this function disagree. The extension clamps *after* an `as usize` cast, so a negative start lands on `signal_len`, the span comes out empty, and the base disappears. Measured against `leech_core` 0.8.0: | map | `leech_core` | leech NumPy | this PR | |---|---|---|---| | `[-8, 10, 20, 30]` | 60 hot | 90 hot | 90 hot | | `[-30, -20, 40, 60]` | 0 hot | 90 hot | 90 hot | | everything else tested | agree | agree | agree | Which one runs today depends on whether `leech_core` is importable in the environment — the divergence is invisible from Python. This crate keeps the surviving tail: it is the readable definition, and it is what a reference-anchored map, whose entries legitimately go negative once the aligned region is cropped, needs. leech's own chunking never produces a negative entry (`chunk_seq_to_sig[0]` is snapped to 0), so no trained model is affected. ## Follow-up, not in this PR Having `leech_core` delegate here, the way it already delegates the refinement preset, the POD5 reader cache, `span_stats` and the mapping primitives. leech pins `escapepod-signal` by git tag, so that waits for the next release — at which point the divergence above disappears rather than being pinned twice. ## Verification - `cargo nextest run --workspace` — 706 passed - `cargo test --doc --workspace` — 26 passed - `cargo clippy --workspace --all-targets` — clean - `cargo fmt --all`
jayhesselberth
added a commit
that referenced
this pull request
Aug 26, 2026
#222) * refactor(rust): take the signal-level k-mer encoding from escapepod-signal leech held the only copy of this rule, inside a `crate-type = ["cdylib"]` Python extension module — so a native runtime for a leech `signal_kmer` model could not link it and had to transcribe it, which is a second definition that diverges silently. That is not hypothetical in this stack: `extract_levels` was written twice with different centring conventions and moved 25 of 100 features, and escapepod-classify reproduced a superseded feature definition for two months, its golden missing it because all 19 fixture reads took the other branch. It is also the natural pair to `escapepod_signal::mapping`, which *produces* the base-to-signal map the encoding consumes. The producing half was already upstream and the consuming half was not. escapepod-rs#272 fixed that; this is leech taking the call (escapepod-signal v0.16.0). Two primitives delegate, both byte-identical to what was here: - `encoding.rs::encode_signal_kmer_inner` -> `seq_encoding::encode_signal_kmer` - `features.rs::sequence_to_int` -> `seq_encoding::sequence_to_int`, the same A/C/G/T-U table with the same `-1` for anything else (upstream's `UNKNOWN_BASE` is `-1`) A third overlap is deliberately NOT delegated. The k-mer context slice in `signal_mapping.rs` covers the same range with the same padding as upstream's `sequence_ints_with_context`, but leech needs it as BASES — it serializes `sequence_with_kmer_context` as a string — where upstream returns ints. Same rule, different type; swapping it would change a serialized chunk field. This is a delegation, so identity is the only acceptable outcome, and the evidence is that nothing had to change to accommodate it: 198 parity tests pass untouched, including `test_backend_parity.py`, which compares every array in the npz between the Rust and Python backends — and the Python reference is not touched by this commit, so it is an independent implementation. Also here, because the pin moves anyway: - `escapepod>=0.16.0`, and the full suite re-run against that Python package as well as against the Rust crate (they are independent: leech_core links the crate statically, `leech.io` imports the wheel). 1467 passed on both. - `uv.lock` gains onnx/onnxruntime/onnxscript/protobuf. #220 added the `onnx` extra to pyproject and never locked it — CI installs from pyproject so it worked, but `uv sync --extra onnx` would not have. - The ONNX contract stops telling consumers to reimplement the encoder and names the crate that owns it. * refactor(rust): take the k-mer context windowing from escapepod-signal too The third and last overlap, and the one that needed an upstream change first (escapepod-rs#274, released in 0.16.1). It could not delegate before: leech needs the context window as BASES, because the corpus serializes `sequence_with_kmer_context` as a string, where upstream only offered ints. escapepod-signal now exposes both forms over one windowing rule, with `sequence_to_int(bases) == ints` pinned by a test there. So all three halves of the signal-level k-mer path live upstream now: `mapping` produces the base-to-signal map, `sequence_bases_with_context` cuts the window it covers plus the k-mer context, and `encode_signal_kmer` scatters the one-hot context along the signal axis. leech held the middle and the last of those, inside a cdylib no Rust consumer could link. This is the one of the three most worth getting out of a second copy. It is where `before` and `after` are not interchangeable — swap them and every k-mer is read from a window displaced by `before - after` bases, silently, and `encode_signal_kmer` cannot detect it because it only sees the total width. Identical by construction and by test: the ranges agree (`(seq_start - before)..(seq_end + after)` is upstream's `core_start - before` plus `n_bases + before + after` with `n_bases = seq_end - seq_start`), and both pad rather than shift, leech with `b'N'` and upstream with `UNKNOWN_BASE_CHAR`, which is `b'N'`. 198 parity tests pass untouched, and `sequence_with_kmer_context` is one of the fields `test_backend_parity.py` compares array-by-array between the Rust and Python backends — the Python side is not touched here, so it is an independent check. Only `rust/Cargo.toml`'s git tag moves to v0.16.1; `Cargo.lock` pins it at b4c9afae. The `escapepod` Python pin stays `>=0.16.0`, because the new function is in the Rust crate and not in the Python bindings, and the two are independent — leech_core links the crate statically, `leech.io` imports the wheel. Full suite 1467 passed, 44 skipped.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #217, and delivers the CRF half of the export at the same time — the two
need the same exporter settings, the same round-trip check, and the same idea of
a contract, so the shared layer is justified by having two consumers rather than
by anticipation.
What lands
leech model export --format onnx, beside the existing--format torch(unchanged default)
leech.crf.export.export_crf_onnx— the encoder only; the decode is notexpressible in standard ONNX ops, which is why
escapepod-demuxowns itleech.onnx_export— the shared exporter, round-trip verifier, and contractonnxextra; CI installs it so the round-trip tests run rather thanskip
The caveat #217 flagged, pinned
dynamo=False— the obvious first attempt — fails on these architectures withUnsupported: ONNX export of operator adaptive_avg_pool1d. That is an exporterlimitation, not a model problem, but the message reads like one. It is
documented where someone will hit it, and
test_export_uses_the_dynamo_pathexports exactly that shape so a regression to the legacy exporter says so.
What a graph cannot carry
Each export writes a contract beside it:
seq_channels = sum(signal_kmer_context) * 4 + 4re-derived by handtarget[state_len:]The consequential decision
#217 asks whether the
signal_kmerencoding belongs in front of the graph orstays a documented call. It stays a documented call, and the contract names
the function and its parameters.
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 bakes
signal_kmer_contextinto the graph. Same choice escapepod-rs's charging bundlemakes: the recipe travels in
metadata.json, never in flags.Correction (1b8f40b)
An earlier revision of this PR said "leech-core already ships the encoder, so
depend on it rather than reimplementing". That is true for a Python consumer and
false for the one that matters:
leech_coreiscrate-type = ["cdylib"], aPython extension module, so escapepod-rs cannot link it. As written it read as
though the one-definition problem were solved. It is not, and the text now says
so.
The real fix is upstream, and is filed as rnabioco/escapepod-rs#271:
encode_signal_kmer_inneris pure, dependency-free and carries no modelvocabulary, and escapepod-signal already owns
mapping— which produces themap this consumes. Producer upstream, consumer downstream, which is backwards.
Until it moves, a Rust consumer transcribes the rule and the mitigation is a
cross-language golden, as with the CRF decode and the charging features.
The decision to keep the encoding out of the graph is unchanged and stands on
its own two reasons above; only the "already solved" framing is gone.
Verification
Across the serialization boundary, not in process:
matching the 4.77e-07 / 1.19e-06 #217 measured on production classifier arms.
Testing
shape, dynamic batch on axis 1 (the CRF graph is time-major, so the naive
axis-0 choice pins the graph to its traced batch), the contract fields, and
the refusal to export without standardisation
ruff format --check,ruff check,ty check,zensical buildclean🤖 Generated with Claude Code