model export: add ONNX beside torch.export — the graphs already convert, with one exporter caveat

[jayhesselberth]

leech model export writes torch.export only (model_export.py:export_single_model). That makes a trained arm loadable by anything with PyTorch, and by nothing else — so a downstream runtime that consumes ONNX cannot run a leech model at all, and a consumer repo ends up recording that the leech arm "is not a shippable format" as if it were a property of the models.

It isn't. I exported two trained production arms and checked them against torch:

arm	export	size	onnxruntime vs torch, max abs diff
ConvLSTMBase	ok	0.4 MB	4.77e-07
TCNDwellResidualLN	ok	0.9 MB	1.19e-06

float32 eps is 1.19e-07, so both are rounding. onnx.checker.check_model passes on both. Real checkpoints at the production geometry (signal_context [90, 300], seq_encoding=signal_kmer, signal_in_channels=2), not toy modules.

The one thing that will bite whoever implements this

dynamo=False fails on both architectures. The legacy TorchScript exporter raises:

SymbolicValueError: Unsupported: ONNX export of operator adaptive_avg_pool1d,
output size that are not factor of input size

and, for the TCN, the input size not accessible variant of the same. That is an exporter limitation, not a property of the model — dynamo=True at opset 18 exports both cleanly. Worth a comment in the code, because the obvious first attempt is the legacy path and its error message reads like a model problem.

Suggested shape

leech model export --format {torch,onnx} (default torch, so nothing changes for existing callers), reusing _build_example_inputs, which already computes the right arity and shapes per architecture:

ConvLSTMBase        (N,2,390) + (N,36,390)
TCNDwellResidualLN  (N,2,390) + (N,36,390) + (N,12,21)

Two things worth writing into the exported artifact, since a consumer cannot recover them from the graph:

1. The input contract — names, shapes, dtypes, and which input is which. config.json has the pieces (signal_len, signal_in_channels, seq_encoding, signal_kmer_context, feature_start/feature_end), but a consumer should not have to re-derive seq_channels = sum(signal_kmer_context) * 4 + 4.
2. The output convention — these emit a single BCE logit (N,1), not a 2-class softmax. Easy to get wrong from the outside.

The part that ONNX alone does not solve

The 36-channel sequence input is encode_signal_kmer output (features.py:482) — a scatter of the one-hot k-mer context along the signal axis through the base-to-signal map. It lives in the dataset, not the model, so it is not in the exported graph, and any non-Python runtime has to produce it before it can call the model.

leech-core already ships that encoder in Rust (_rs_encode_signal_kmer). A consumer depending on it keeps one definition of the rule; reimplementing it creates a second, and the two diverge silently — a downstream repo here had exactly that happen, reproducing a superseded feature definition for two months before a real-corpus comparison caught it in 4 reads out of 842.

So it may be worth deciding whether the encoding belongs in front of the graph (as an ONNX prefix taking signal + map) or stays a documented leech-core call. That choice is more consequential than the export itself.

Repro

escapepod-models has a probe at scripts/charging/probe_onnx_export.py that exports a checkpoint dir, runs onnx.checker, and compares onnxruntime against torch on identical inputs. Points at any config.json + model_best.pt directory.