Skip to content

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

Merged
jayhesselberth merged 1 commit into
mainfrom
fix-adaptive-pool-export
Aug 30, 2026
Merged

fix(export): make the ONNX loadable by a non-Python runtime#233
jayhesselberth merged 1 commit into
mainfrom
fix-adaptive-pool-export

Conversation

@jayhesselberth

Copy link
Copy Markdown
Member

Unblocks rnabioco/escapepod-models#96. charging_tcn_rna004@v0.1.0 shipped a
graph no released escpod can load; this is the export-side fix, same weights,
no retrain.

Note the diagnosis in escapepod-models#96 is wrong — it names
nn.MultiheadAttention, which in fact exports as plain MatMul/Softmax/MatMul.
The culprit is adaptive_avg_pool1d. escapepod-rs's waveform_net.rs docs
carry the same wrong attribution and want correcting separately.


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.


What lands after this

  1. Release leech (minor — the pool's fp rounding moves ~2.4e-07).
  2. Re-pin leech in escapepod-models' pixi.toml.
  3. Re-export and ship charging_tcn_rna004@v0.1.1.

The registry-side gate that would have caught this at build time is a separate
PR in escapepod-models and lands independently of this one.

Not verified

That a released escpod runs the resulting bundle end to end — 0.18.1 has no
waveform bundle variant, so "tract loads and runs the graph" is the extent of
the evidence here.

`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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017QSUHQ2x8ZGh9q5GoY8mw4
@jayhesselberth
jayhesselberth merged commit 3d20bb7 into main Aug 30, 2026
3 checks passed
@jayhesselberth
jayhesselberth deleted the fix-adaptive-pool-export branch August 30, 2026 23:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant