From 2add3e273b9c79b6a1a5275bbacdcf1e0da75a2e Mon Sep 17 00:00:00 2001 From: Jay Hesselberth Date: Sun, 30 Aug 2026 17:05:20 -0600 Subject: [PATCH] fix(export): make the ONNX loadable by a non-Python runtime `charging_tcn_rna004@v0.1.0` shipped from this exporter with a graph no released `escpod` binary can load: tract parses it and then gives up during shape analysis, five ways (rnabioco/escapepod-models#96). onnxruntime loads it fine, which is why `verify_onnx` had nothing to say and the failure surfaced at integration instead of at build time. Two independent causes, both measured against tract 0.23.5 through the load path `escapepod_classify` actually uses (pin the batch with `with_input_fact`, rewrite nothing else in the proto): 1. `adaptive_avg_pool1d` with an output size that does not divide the input (390 -> 11 here). Dynamo open-codes it as `Unsqueeze -> Transpose -> GatherND -> Transpose -> Where(masked_fill)` plus one `Gather`/`Add` per element of the widest bin: a rank-8 gather over an all-constant index and mask. tract fails on it pinned (`Val(64) vs Val(1)`), unpinned (`Sym(batch) vs Val(1)`), and with `value_info` cleared it dies one node later on the rank-8 `Transpose`. No post-hoc rewrite helps -- onnx-simplifier folds away every `Shape` node and leaves the `GatherND`; onnxruntime's optimiser keeps it and adds ORT-only fusions. `models.components.AdaptiveAvgPool1d` now writes the same arithmetic as one matmul against a constant `[L_in, L_out]` segment-mean matrix. The bin rule is PyTorch's own, `[floor(j*L/K), ceil((j+1)*L/K))`, upsampling included -- `ResNetDwell` pools 4 up to 11. Agreement with the aten op is 2.4e-07 over a grid of lengths and output sizes, and the matmul runs in float32 outside autocast so the accumulation matches what the aten op does under AMP. One implementation, so `nn.AdaptiveAvgPool1d` (the registry layer), `resnet_dwell`, `transformer_dwell` and the `tests/reference_*` oracles all move together and the bit-exact config-vs-reference parity tests stay bit-exact. `signal_cnn`'s `AdaptiveAvgPool1d(1)` is left alone: 1 divides everything, it exports as `GlobalAveragePool`, and pinning its length would be a regression. 2. `value_info`. Dynamo writes one entry per intermediate -- 667 for this model -- with the batch axis as the *symbol* `batch`, because that is what `dynamic_axes` asked for. A consumer that pins the batch then cannot unify, and tract fails at the FIRST convolution: Failed analyse for node "node_conv1d" ConvHir: Unifying shapes batch,64,390 and 1,64,390: Impossible to unify Sym(batch) with Val(1) `strip_value_info` drops them and `export_onnx` always calls it. Nothing needs them: every runtime re-infers, `onnx.checker` is satisfied, and every graph escpod loads today has zero -- the legacy exporter never wrote any, which is why its graphs always loaded. Initializers are untouched, external data references included. Measured on the shipped `TCNDwellResidualLN` weights, no retrain: nodes 479 -> 319, GatherND 2 -> 0, Gather 76 -> 0 tract, batch 1 and 32 loads, optimizes and runs (was: five failures) tract vs torch max |dlogit| 5.72e-06 over 256 real chunks, 0 decision disagreements onnxruntime vs torch 1.335e-05 over 4096 real chunks (shipped: 1.4305e-05) The module docstring's case for the dynamo exporter was re-measured rather than inherited, since the pool no longer emits an aten adaptive pool and that could have retired it. It did not: `dynamo=False` still refuses both the aten pool and leech's replacement, because `torch.jit.trace` turns `.shape[-1]` into a Tensor and takes the dynamic-length fallback. The docstring now says so, and a test pins it. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_017QSUHQ2x8ZGh9q5GoY8mw4 --- src/leech/models/components.py | 122 +++++++++++++++++++++++++ src/leech/models/nn.py | 5 +- src/leech/models/resnet_dwell.py | 6 +- src/leech/models/transformer_dwell.py | 6 +- src/leech/onnx_export.py | 74 ++++++++++++++- tests/reference_conv_lstm.py | 16 ++-- tests/reference_tcn.py | 16 ++-- tests/test_onnx_export.py | 124 ++++++++++++++++++++++++++ 8 files changed, 348 insertions(+), 21 deletions(-) diff --git a/src/leech/models/components.py b/src/leech/models/components.py index b28a057..8de8984 100644 --- a/src/leech/models/components.py +++ b/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. diff --git a/src/leech/models/nn.py b/src/leech/models/nn.py index 276db45..8c6a4ee 100644 --- a/src/leech/models/nn.py +++ b/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} diff --git a/src/leech/models/resnet_dwell.py b/src/leech/models/resnet_dwell.py index 0db6de9..9153133 100644 --- a/src/leech/models/resnet_dwell.py +++ b/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( diff --git a/src/leech/models/transformer_dwell.py b/src/leech/models/transformer_dwell.py index e65a682..7218533 100644 --- a/src/leech/models/transformer_dwell.py +++ b/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) diff --git a/src/leech/onnx_export.py b/src/leech/onnx_export.py index cbb59c4..ae24da8 100644 --- a/src/leech/onnx_export.py +++ b/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, diff --git a/tests/reference_conv_lstm.py b/tests/reference_conv_lstm.py index 243710f..01be147 100644 --- a/tests/reference_conv_lstm.py +++ b/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( diff --git a/tests/reference_tcn.py b/tests/reference_tcn.py index 0e5a15b..81e040d 100644 --- a/tests/reference_tcn.py +++ b/tests/reference_tcn.py @@ -35,7 +35,7 @@ DEFAULT_SIGNAL_KMER_CONTEXT, DEFAULT_SIGNAL_LEN, ) -from leech.models.components import BaseModel, FeatureBranch, make_norm +from leech.models.components import AdaptiveAvgPool1d, BaseModel, FeatureBranch, make_norm # Default: first 5 channels are dwell features NUM_DWELL_FEATURES = 5 @@ -237,7 +237,7 @@ def __init__( dropout=dropout, norm_type=norm_type, ) - self.signal_pool = nn.AdaptiveAvgPool1d(kmer_len) + self.signal_pool = AdaptiveAvgPool1d(kmer_len) # Sequence branch: TCN self.seq_tcn = TCN( @@ -249,7 +249,7 @@ def __init__( norm_type=norm_type, ) 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 self.feature_branch = FeatureBranch( @@ -380,7 +380,7 @@ def __init__( dropout=dropout, norm_type=norm_type, ) - self.signal_pool = nn.AdaptiveAvgPool1d(kmer_len) + self.signal_pool = AdaptiveAvgPool1d(kmer_len) # Sequence branch: TCN self.seq_tcn = TCN( @@ -392,7 +392,7 @@ def __init__( norm_type=norm_type, ) 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 self.feature_branch = FeatureBranch( @@ -513,7 +513,7 @@ def __init__( dropout=dropout, norm_type=norm_type, ) - self.signal_pool = nn.AdaptiveAvgPool1d(kmer_len) + self.signal_pool = AdaptiveAvgPool1d(kmer_len) self.residual_tcn = TCN( in_channels=1, @@ -523,7 +523,7 @@ def __init__( dropout=dropout, norm_type=norm_type, ) - self.residual_pool = nn.AdaptiveAvgPool1d(kmer_len) + self.residual_pool = AdaptiveAvgPool1d(kmer_len) self.seq_tcn = TCN( in_channels=seq_in_channels, @@ -534,7 +534,7 @@ def __init__( norm_type=norm_type, ) if seq_encoding == "signal_kmer": - self.seq_pool = nn.AdaptiveAvgPool1d(kmer_len) + self.seq_pool = AdaptiveAvgPool1d(kmer_len) self.feature_branch = FeatureBranch( num_features=num_features, conv_channels=conv_channels, norm_type=norm_type diff --git a/tests/test_onnx_export.py b/tests/test_onnx_export.py index 29d3640..5683f6b 100644 --- a/tests/test_onnx_export.py +++ b/tests/test_onnx_export.py @@ -269,6 +269,130 @@ def forward(self, x): assert verify_onnx(path, model, example, input_names=["x"]) < 1e-5 +def test_the_legacy_exporter_still_refuses_the_aten_adaptive_pool(tmp_path): + """The claim the module docstring rests on, measured rather than inherited. + + It matters that this is checked and not assumed: leech's own + `AdaptiveAvgPool1d` no longer emits an aten adaptive pool, which could + plausibly have retired the reason for pinning the dynamo exporter. It did + not — `torch.jit.trace` makes `.shape[-1]` a Tensor, so the replacement + takes its dynamic-length fallback and lands on the same aten op. + """ + from leech.models.components import AdaptiveAvgPool1d + + example = (torch.randn(2, 3, 100),) + for name, pool in (("aten", torch.nn.AdaptiveAvgPool1d(7)), ("leech", AdaptiveAvgPool1d(7))): + model = torch.nn.Sequential(pool, torch.nn.Flatten(1)).eval() + with pytest.raises(Exception, match="adaptive_avg_pool1d"): + torch.onnx.export( + model, + example, + str(tmp_path / f"{name}.onnx"), + dynamo=False, + opset_version=OPSET, + input_names=["x"], + output_names=["y"], + ) + + +# ── what a non-Python runtime needs, and onnxruntime never notices ───────── + + +def test_the_pool_exports_as_a_matmul_not_a_gather(tmp_path): + """A non-dividing adaptive pool must not reach the graph as `GatherND`. + + 390 -> 11 is the charging TCN's geometry. Dynamo open-codes the aten op as + `Unsqueeze -> Transpose -> GatherND -> Transpose -> Where`, a rank-8 gather + over an all-constant index that tract 0.23.5 cannot close — which is what + made `charging_tcn_rna004@v0.1.0` unloadable by every released `escpod` + (rnabioco/escapepod-models#96). onnxruntime runs it fine, so no round-trip + check can see this; only the op list can. + """ + import onnx + + from leech.models.components import AdaptiveAvgPool1d + + model = torch.nn.Sequential(AdaptiveAvgPool1d(11), torch.nn.Flatten(1)).eval() + example = (torch.randn(2, 3, 390),) + path = export_onnx( + model, example, tmp_path / "pool.onnx", input_names=["x"], output_names=["y"] + ) + ops = {n.op_type for n in onnx.load(str(path)).graph.node} + assert "GatherND" not in ops + assert "MatMul" in ops + assert verify_onnx(path, model, example, input_names=["x"]) < 1e-5 + + +def test_the_pool_matches_the_aten_op_it_replaces(): + """Same arithmetic, to float32 rounding, across the bin-width edge cases. + + `adaptive_avg_pool1d`'s bins are `[floor(j*L/K), ceil((j+1)*L/K))`, so + widths differ by one whenever K does not divide L — 390 -> 11 gives bins of + 36 and 37. A matrix built on a different convention would agree on the + dividing cases and quietly disagree on exactly the ones that matter. + + The grid includes `k > length`, which is not a degenerate case anyone should + skip: `ResNetDwell` pools a length-4 map up to 11, so the op UPSAMPLES + there. A range guard that rejected it shipped in the first draft of this + and was caught by two model tests rather than by this one; the grid now + covers it. + """ + from leech.models.components import AdaptiveAvgPool1d + + torch.manual_seed(0) + worst = 0.0 + for length in (3, 4, 11, 12, 37, 100, 390, 400, 1024): + for k in (1, 5, 7, 11, 21): + x = torch.randn(3, 8, length) + want = torch.nn.functional.adaptive_avg_pool1d(x, k) + got = AdaptiveAvgPool1d(k)(x) + assert got.shape == want.shape + worst = max(worst, float((want - got).abs().max())) + assert worst < 10 * EPS, worst + + +def test_export_writes_no_value_info(tmp_path): + """Dynamo records every intermediate's shape with the batch axis as a + SYMBOL; a consumer that pins the batch then cannot unify, and tract fails + at the first convolution with `Sym(batch) vs Val(1)`. Every graph escpod + loads carries zero entries.""" + import onnx + + model = torch.nn.Sequential(torch.nn.Conv1d(3, 4, 3), torch.nn.Flatten(1)).eval() + path = export_onnx( + model, + (torch.randn(2, 3, 32),), + tmp_path / "conv.onnx", + input_names=["x"], + output_names=["y"], + ) + assert list(onnx.load(str(path)).graph.value_info) == [] + + +def test_strip_value_info_is_idempotent_and_reports(tmp_path): + """It returns how many it dropped, so a caller can log a real number, and + running it twice is not an error.""" + import onnx + + from leech.onnx_export import strip_value_info + + model = torch.nn.Sequential(torch.nn.Conv1d(3, 4, 3), torch.nn.Flatten(1)).eval() + path = export_onnx( + model, + (torch.randn(2, 3, 32),), + tmp_path / "again.onnx", + input_names=["x"], + output_names=["y"], + ) + # export_onnx already stripped, so there is nothing left to drop. + assert strip_value_info(path) == 0 + proto = onnx.load(str(path)) + proto.graph.value_info.extend(proto.graph.input) + onnx.save(proto, str(path)) + assert strip_value_info(path) == len(proto.graph.input) + assert list(onnx.load(str(path)).graph.value_info) == [] + + def test_verify_returns_the_actual_difference(tmp_path): model = torch.nn.Linear(4, 2).eval() example = (torch.randn(3, 4),)