Skip to content

docs(classify): tract dies on the adaptive pool, not on attention - #309

Merged
jayhesselberth merged 1 commit into
worktree-issue-306-waveform-chunksfrom
fix-waveform-attribution
Aug 31, 2026
Merged

docs(classify): tract dies on the adaptive pool, not on attention#309
jayhesselberth merged 1 commit into
worktree-issue-306-waveform-chunksfrom
fix-waveform-attribution

Conversation

@jayhesselberth

Copy link
Copy Markdown
Member

Docs only, no code changes. Targets worktree-issue-306-waveform-chunks, not
main, because that is where waveform_net.rs lives.

waveform_net.rs attributes the unloadable charging_tcn_rna004@v0.1.0 graph
to dynamo's lowering of nn.MultiheadAttention's mask handling, and explicitly
rules out adaptive_avg_pool1d ("those layers are in the model config and are
not what tract dies on"). It is the other way round. nn.MultiheadAttention
exports as plain MatMul/Softmax/MatMul; the culprit is
adaptive_avg_pool1d with an output size that does not divide the input.

The evidence is already in the file

Every measurement here was right — the five failure modes, the node names, and
the two constant initializers. The constants settle it:

charging_tcn_rna004 pools 390 down to 11. PyTorch's bin rule is
[floor(j*L/K), ceil((j+1)*L/K)), which for 390 -> 11 gives eleven bins of
width 36 or 37. A gather that evaluates every bin at once therefore needs
an (11, 37) index grid and an (11, 37) mask marking the slot the 36-wide
bins do not use — which is exactly the (1,1,11,37,2) index tensor and
(11,37) bool mask this file records. An attention mask is shaped by sequence
length and head count and would never be 11 x 37.

Why it was easy to get wrong

A non-dividing adaptive pool lowers to no ONNX pooling op at all. Grepping
the graph for one finds nothing, so the ragged-bin gather looks like it must
have come from somewhere else. The layers named in escapepod-rs#306 were the
right suspects — they just do not appear under a pooling name.

The second cause, which the table shows but the prose never named

The value_info dynamo writes for every intermediate carries the batch axis as
the symbol batch. That is the node_conv1d row: a consumer that pins the
batch cannot unify against a symbol, so it fails at the first convolution,
before ever reaching the gather. Two independent causes, not one.

Upstream status

Diagnosed and fixed in rnabioco/leech#233, released in leech 0.10.0: the
pool becomes one matmul against a constant segment-mean matrix, and every
export strips value_info. Measured there 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 flips

Neither cause is visible to a round-trip check against onnxruntime, which loads
the old graph happily — which is why this surfaced at integration here rather
than at build time in leech.

So the ort bridge can be removed once charging_tcn_rna004 is re-exported
and shipped (rnabioco/escapepod-models#96). Still unverified: a released
escpod running such a bundle end to end, since 0.18.1 has no waveform bundle
variant.

Scope

Corrected in all five places the claim had spread to — waveform_net.rs,
CHANGELOG.md, CLAUDE.md, docs/cli/signal-classify.md, and
examples/tract_dynamo_probe.rs — since leaving it standing in the changelog
and the agent instructions would be no fix at all.

cargo fmt --all --check and cargo clippy --workspace --all-targets --features escapepod-classify/waveform-onnx both clean.

This repo diagnosed the unloadable `charging_tcn_rna004@v0.1.0` graph as
dynamo's lowering of `nn.MultiheadAttention`'s mask handling, and
`waveform_net.rs` explicitly ruled out `adaptive_avg_pool1d`. That is
backwards. MHA exports as plain MatMul/Softmax/MatMul; the culprit is
`adaptive_avg_pool1d` with an output size that does not divide the
input, which dynamo open-codes as the ragged-bin gather this file
already describes.

The measurements here were always right -- the five failure modes, the
node names, and the two constants. It is the constants that settle it:
`charging_tcn_rna004` pools 390 down to 11, and PyTorch's bin rule
`[floor(j*L/K), ceil((j+1)*L/K))` gives eleven bins of width 36 or 37,
so evaluating every bin at once needs exactly the `(1,1,11,37,2)` index
tensor and `(11,37)` mask recorded above. An attention mask is shaped by
sequence length and head count and is never `11 x 37`.

What made this easy to get wrong is worth keeping: a non-dividing
adaptive pool lowers to no ONNX pooling op at all, so grepping the graph
for one finds nothing and the gather looks like it came from elsewhere.
The layers named in the model config were the right suspects; they just
do not appear under a pooling name. escapepod-rs#306 had it right.

Also records the second cause the table here shows but the prose did not
name: the symbolic `value_info` dynamo writes per intermediate, which is
why the pinned-batch run fails at the *first* convolution rather than at
the gather.

Fixed upstream in rnabioco/leech#233, released in leech 0.10.0 -- the
pool becomes one matmul against a constant segment-mean matrix and every
export strips `value_info`. On the shipped weights that is 479 -> 319
nodes, GatherND 2 -> 0, and tract loads, optimizes and runs the graph
within 5.72e-06 of torch. So the `ort` bridge can go once
`charging_tcn_rna004` is re-exported (rnabioco/escapepod-models#96);
what remains unverified is a released escpod running such a bundle end
to end, since 0.18.1 ships no waveform variant.

Docs only -- no code changes. Corrected in all five places the claim had
spread to, since leaving it standing in the CHANGELOG and CLAUDE.md
would be no fix at all.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GK2ESAWBrvVbsqG3vbqDoT
@jayhesselberth
jayhesselberth merged commit e4880cb into worktree-issue-306-waveform-chunks Aug 31, 2026
@jayhesselberth
jayhesselberth deleted the fix-waveform-attribution branch August 31, 2026 13:12
jayhesselberth added a commit that referenced this pull request Sep 1, 2026
#309 landed the same correction this branch reached independently — that tract
dies on `adaptive_avg_pool1d`, not on `nn.MultiheadAttention` — so every
conflict is two descriptions of one finding.

Kept #309's identification, which is better than mine: it derives the `(11, 37)`
index grid and mask from PyTorch's bin rule (390 -> 11 gives eleven bins of
width 36 or 37) instead of asserting the shapes, names the trap that a
non-dividing adaptive pool lowers to no ONNX pooling op at all so grepping for
one finds nothing, and records that neither cause is visible to a round-trip
against onnxruntime — which is why this surfaced at integration.

Kept this branch's runtime prose everywhere the two disagree about what the code
does, because the code moved underneath #309: the graph now runs through tract
against the `@v0.1.1` re-export, so `classify-waveform`, `ORT_DYLIB_PATH` and
the "cannot go through tract" framing no longer describe anything. #309's
`docs/cli/signal-classify.md` edit lands on `docs/cli/classify.md`, which is
where #307 moved that page.
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