fix(export): make the ONNX loadable by a non-Python runtime

src/leech/models/components.py

@@ -5,6 +5,9 @@
 eliminating code duplication and making it easier to create new models.
 """
 
+import functools
+
+import numpy as np
 import torch
 import torch.nn as nn
 
@@ -17,6 +20,125 @@
 )
 
 
+@functools.lru_cache(maxsize=128)
+def _segment_mean_weights(length: int, output_size: int) -> np.ndarray:
+    """The cached half of :func:`segment_mean_matrix`, as **numpy**.
+
+    Deliberately not a ``torch.Tensor``. Under ``torch.export``'s fake-tensor
+    tracing, ``torch.from_numpy`` returns a *FakeTensor*; cached, that fake
+    outlives the trace and poisons every later eager call — the model then
+    returns a tensor subclass and ``verify_onnx``'s ``.numpy()`` dies a long
+    way from the cause. Cache the array, build the tensor per call: the build
+    is a 4 KB copy and the bin arithmetic is what was worth caching.
+
+    Built in float64 so the reciprocals are exact before they are rounded once
+    into the model's dtype.
+
+    ``output_size > length`` is legal and is not a mistake: ``ResNetDwell``
+    pools a length-4 feature map up to ``kmer_len`` 11, and ``adaptive_avg_pool``
+    UPSAMPLES there by repeating bins. The bin formula covers it unchanged —
+    ``ceil((j+1)*L/K) > floor(j*L/K)`` for every ``j`` whenever ``L >= 1``, so no
+    bin is ever empty. An earlier version of this rejected it as out of range
+    and two model tests caught it.
+    """
+    if length < 1 or output_size < 1:
+        raise ValueError(f"length {length} and output_size {output_size} must both be >= 1")
+    w = np.zeros((length, output_size), dtype=np.float64)
+    for j in range(output_size):
+        start = (j * length) // output_size
+        end = -((-(j + 1) * length) // output_size)  # ceil((j+1)*L/K)
+        w[start:end, j] = 1.0 / (end - start)
+    return w
+
+
+def segment_mean_matrix(
+    length: int,
+    output_size: int,
+    *,
+    dtype: torch.dtype = torch.float32,
+    device: torch.device | str | None = None,
+) -> torch.Tensor:
+    """``[length, output_size]``: column *j* averages ``x[start_j:end_j]``.
+
+    PyTorch's adaptive-pool bin rule, written out: bin *j* covers
+    ``[floor(j*L/K), ceil((j+1)*L/K))``, so the bins tile the axis and their
+    widths differ by at most one. Right-multiplying by this matrix *is*
+    ``adaptive_avg_pool1d`` — see :class:`AdaptiveAvgPool1d` for why leech
+    spells it that way.
+    """
+    return torch.tensor(
+        _segment_mean_weights(int(length), int(output_size)), dtype=dtype, device=device
+    )
+
+
+class AdaptiveAvgPool1d(nn.AdaptiveAvgPool1d):
+    """``nn.AdaptiveAvgPool1d`` written as one matmul against a constant.
+
+    Same arithmetic, same ``output_size``, no parameters — and an ONNX graph a
+    non-Python runtime can actually load. ``aten::adaptive_avg_pool1d`` has no
+    ONNX op behind it when the output size does not divide the input size, so
+    every exporter has to open-code it:
+
+    * the **legacy TorchScript** exporter refuses outright
+      (``SymbolicValueError: ... output size that are not factor of input
+      size``);
+    * the **dynamo** exporter open-codes it as
+      ``Unsqueeze -> Transpose -> GatherND -> Transpose -> Where(masked_fill)``
+      followed by one ``Gather``+``Add`` per element of the widest bin. That is
+      a rank-8 gather over an all-constant index and mask, and tract 0.23.5
+      gives up on it during shape analysis — which is what made
+      ``charging_tcn_rna004@v0.1.0`` unloadable by the shipped ``escpod``
+      binaries (rnabioco/escapepod-models#96). No graph rewrite fixes it
+      afterwards: onnx-simplifier folds away every ``Shape`` node and still
+      leaves the ``GatherND``, and onnxruntime's optimiser keeps it and adds
+      ORT-only fusions on top.
+
+    Written out as a matmul, the same two pools become one ``MatMul`` each
+    against a ``[L_in, L_out]`` initializer, and the graph loads.
+
+    The matmul is forced to float32 outside autocast: ``adaptive_avg_pool1d``
+    is not on autocast's cast list, so it keeps its input dtype and accumulates
+    in float32, and a bare ``@`` under autocast would silently become an fp16
+    gemm. Keeping the accumulation in float32 is what makes this a rewrite of
+    the same function rather than a change to it.
+    """
+
+    def _output_size(self) -> int:
+        size = self.output_size
+        if isinstance(size, tuple | list):
+            (size,) = size
+        return int(size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        length = x.shape[-1]
+        if not isinstance(length, int):
+            # The pooled axis is not a concrete length, so the bin edges are
+            # not knowable and there is no constant matrix to build. Fall back
+            # to the aten op — which is the graph this class exists to avoid,
+            # so it is worth knowing exactly when this fires:
+            #
+            #   eager                     never (a shape is an int)
+            #   torch.export / dynamo     only if the LAST axis is declared
+            #                             dynamic; leech's exports make only
+            #                             the batch axis dynamic, and the
+            #                             charging TCN's graph is 8 MatMuls
+            #                             and zero GatherND, so it does not
+            #   torch.compile, dynamic    yes, and correctly — an unknown
+            #                             length has no constant matrix
+            #   torch.jit.trace           YES. Tracing makes `.shape[-1]` a
+            #                             Tensor, so this branch is taken and
+            #                             the legacy TorchScript ONNX
+            #                             exporter still meets an
+            #                             `adaptive_avg_pool1d` it refuses.
+            #                             Measured, not assumed; see
+            #                             `leech.onnx_export`'s docstring.
+            return nn.functional.adaptive_avg_pool1d(x, self._output_size())
+        w = segment_mean_matrix(length, self._output_size(), dtype=torch.float32, device=x.device)
+        with torch.amp.autocast(x.device.type, enabled=False):
+            out = torch.matmul(x.float(), w)
+        return out.to(x.dtype)
+
+
 def make_norm(norm_type: str, num_channels: int) -> nn.Module:
     """Create a normalization layer.
 

src/leech/models/nn.py

@@ -177,7 +177,10 @@ def to_dict(self, include_weights: bool = False) -> dict:
 
 
 @register
-class AdaptiveAvgPool1d(nn.AdaptiveAvgPool1d):
+class AdaptiveAvgPool1d(components.AdaptiveAvgPool1d):
+    """See :class:`leech.models.components.AdaptiveAvgPool1d`: the arithmetic of
+    ``nn.AdaptiveAvgPool1d``, written as a matmul so the ONNX graph loads."""
+
     def to_dict(self, include_weights: bool = False) -> dict:
         return {"output_size": self.output_size}
 

src/leech/models/resnet_dwell.py

@@ -25,7 +25,7 @@
     DEFAULT_SIGNAL_KMER_CONTEXT,
     DEFAULT_SIGNAL_LEN,
 )
-from leech.models.components import BaseModel, FeatureBranch
+from leech.models.components import AdaptiveAvgPool1d, BaseModel, FeatureBranch
 
 
 class ResidualBlock1D(nn.Module):
@@ -229,8 +229,8 @@ def __init__(
         # Pool ResNet outputs to kmer_len positions for cross-attention Q
         signal_final_ch = self.signal_resnet.final_channels
         seq_final_ch = self.seq_resnet.final_channels
-        self.signal_pool = nn.AdaptiveAvgPool1d(kmer_len)
-        self.seq_pool = nn.AdaptiveAvgPool1d(kmer_len)
+        self.signal_pool = AdaptiveAvgPool1d(kmer_len)
+        self.seq_pool = AdaptiveAvgPool1d(kmer_len)
 
         # Feature branch: Conv1d on full-width features for cross-attention K/V
         self.feature_branch = FeatureBranch(

src/leech/models/transformer_dwell.py

@@ -27,7 +27,7 @@
     DEFAULT_SIGNAL_KMER_CONTEXT,
     DEFAULT_SIGNAL_LEN,
 )
-from leech.models.components import BaseModel
+from leech.models.components import AdaptiveAvgPool1d, BaseModel
 
 
 class PositionalEncoding(nn.Module):
@@ -133,7 +133,7 @@ def __init__(
             ),
             nn.ReLU(),
         )
-        self.signal_pool = nn.AdaptiveAvgPool1d(kmer_len)
+        self.signal_pool = AdaptiveAvgPool1d(kmer_len)
 
         # Sequence branch: Project encoding to d_model dimensions
         self.seq_conv = nn.Sequential(
@@ -143,7 +143,7 @@ def __init__(
             nn.ReLU(),
         )
         if seq_encoding == "signal_kmer":
-            self.seq_pool = nn.AdaptiveAvgPool1d(kmer_len)
+            self.seq_pool = AdaptiveAvgPool1d(kmer_len)
 
         # Feature branch: Conv1d on full-width features for cross-attention K/V
         # Input: (batch, num_features, kmer_len + margin_left + margin_right)

src/leech/onnx_export.py

@@ -17,6 +17,32 @@
 one and will send whoever hits it looking in the wrong place. ``dynamo=True`` at
 opset 18 exports the same modules cleanly.
 
+Re-measured 2026-08-30, on torch 2.12, because
+:class:`leech.models.components.AdaptiveAvgPool1d` no longer emits an aten
+adaptive pool and that could have retired the reason. **It did not.** Both
+``nn.AdaptiveAvgPool1d(7)`` and leech's replacement, on a length-100 input,
+still fail ``dynamo=False`` with the message above — the replacement because
+``torch.jit.trace`` turns ``.shape[-1]`` into a Tensor, which takes its
+dynamic-length fallback straight back to the aten op. The paragraph stands as
+written; it is now measured rather than inherited.
+
+**But dynamo alone is not sufficient**, which is the half that was missing and
+cost `charging_tcn_rna004@v0.1.0` a release nothing could run
+(rnabioco/escapepod-models#96). Two of its graphs' properties were unloadable
+by tract, the ONNX runtime `escpod` links statically:
+
+* ``adaptive_avg_pool1d`` with a non-dividing output size, which dynamo
+  open-codes as a rank-8 ``GatherND``. Fixed in the model, by
+  :class:`leech.models.components.AdaptiveAvgPool1d`.
+* the ``value_info`` dynamo writes for every intermediate, carrying the batch
+  axis as a symbol. Fixed here, by :func:`strip_value_info`, which
+  :func:`export_onnx` now always calls.
+
+Neither is fixable in the consumer, and neither shows up in
+:func:`verify_onnx` — onnxruntime loads both graphs happily. "It exports and
+round-trips" is a weaker claim than "a runtime can load it", and only the
+second one ships.
+
 What a graph cannot carry
 -------------------------
 Two things a consumer needs and cannot recover from the ONNX file, so both are
@@ -50,6 +76,7 @@
     "contract",
     "describe_inputs",
     "export_onnx",
+    "strip_value_info",
     "verify_onnx",
 ]
 
@@ -151,10 +178,55 @@ def export_onnx(
         output_names=output_names,
         dynamic_axes=dynamic_axes,
     )
-    logger.info("wrote %s (%.2f MB, opset %d)", path, path.stat().st_size / 1e6, opset)
+    dropped = strip_value_info(path)
+    logger.info(
+        "wrote %s (%.2f MB, opset %d, %d value_info entries dropped)",
+        path,
+        path.stat().st_size / 1e6,
+        opset,
+        dropped,
+    )
     return path
 
 
+def strip_value_info(path: str | Path) -> int:
+    """Drop the graph's inferred intermediate shapes. Returns how many went.
+
+    ``value_info`` is optional: it records the shapes the *exporter* inferred
+    for intermediate tensors, and every runtime re-infers them anyway. The
+    dynamo exporter writes one entry per intermediate — 667 of them for the
+    charging TCN — and writes the batch axis into them as the **symbol**
+    ``batch``, because that is what ``dynamic_axes`` asked for.
+
+    That is a contradiction waiting for a consumer that pins the batch.
+    ``escapepod_classify`` loads every graph with
+    ``with_input_fact(0, f32::fact([1, ...]))``, and tract then has to unify a
+    declared ``Sym(batch)`` with the pinned ``Val(1)``, which it cannot:
+
+        Failed analyse for node "node_conv1d" ConvHir: Unifying shapes
+        batch,64,390 and 1,64,390: Impossible to unify Sym(batch) with Val(1)
+
+    — at the *first convolution*, nowhere near anything interesting. Removing
+    the entries lets tract infer from the pinned input and the graph loads.
+    Nothing needs them: onnxruntime re-infers, ``onnx.checker`` is satisfied,
+    and the legacy TorchScript exporter never wrote them in the first place,
+    which is why the graphs it produced (``charging_feature_nn_rna004@v0.1.0``)
+    always loaded. See rnabioco/escapepod-models#96.
+
+    Initializers are left exactly as they are, external-data references
+    included: the proto is read without resolving them and written straight
+    back, so an ``.onnx.data`` sidecar keeps working.
+    """
+    import onnx
+
+    path = Path(path)
+    proto = onnx.load(str(path), load_external_data=False)
+    n = len(proto.graph.value_info)
+    del proto.graph.value_info[:]
+    onnx.save(proto, str(path))
+    return n
+
+
 def verify_onnx(
     path: str | Path,
     model,

tests/reference_conv_lstm.py

@@ -31,7 +31,13 @@
     DEFAULT_SIGNAL_KMER_CONTEXT,
     DEFAULT_SIGNAL_LEN,
 )
-from leech.models.components import BaseModel, FeatureBranch, SequenceBranch, SignalBranch
+from leech.models.components import (
+    AdaptiveAvgPool1d,
+    BaseModel,
+    FeatureBranch,
+    SequenceBranch,
+    SignalBranch,
+)
 
 
 class _ConvLSTMNoAttn(BaseModel):
@@ -90,9 +96,9 @@ def __init__(
                 num_features=num_features, conv_channels=conv_channels, **norm_kw
             )
 
-        self.signal_pool = nn.AdaptiveAvgPool1d(kmer_len)
+        self.signal_pool = AdaptiveAvgPool1d(kmer_len)
         if seq_encoding == "signal_kmer":
-            self.seq_pool = nn.AdaptiveAvgPool1d(kmer_len)
+            self.seq_pool = AdaptiveAvgPool1d(kmer_len)
 
         n_branches = 3 if _has_features else 2
         self.lstm = nn.LSTM(
@@ -240,9 +246,9 @@ def __init__(
                 num_features=num_features, conv_channels=conv_channels, **norm_kw
             )
 
-        self.signal_pool = nn.AdaptiveAvgPool1d(kmer_len)
+        self.signal_pool = AdaptiveAvgPool1d(kmer_len)
         if seq_encoding == "signal_kmer":
-            self.seq_pool = nn.AdaptiveAvgPool1d(kmer_len)
+            self.seq_pool = AdaptiveAvgPool1d(kmer_len)
 
         # BiLSTM on merged signal + sequence
         self.lstm = nn.LSTM(