Unify training and ONNX-export environment, lift PyTorch 1.13 pin

[KakaruHayate]

The historical PyTorch 1.13 pin existed because two bugs surfaced in newer torch when exporting to ONNX. Both are now fixed; training and export can share a single environment (PyTorch >= 2.4).

Bug 1 - WaveNet diffusion ONNX export crashes on PyTorch >= 2.0
modules/backbones/wavenet.py:83 used spec.squeeze(1), which the ONNX
tracer lowers to an onnx::If whose two branches have different ranks
(block0 Squeeze -> rank-3, block1 Identity -> rank-4). Shape inference
for the downstream Conv then fails with SymbolicValueError. Replaced
with spec[:, 0] - an unconditional rank-reducing gather, semantically
identical (eager max-diff = 0) and producing a clean ONNX graph.

Bug 2 - non-RoPE encoder ONNX inference fails at dynamic lengths
torch.nn.MultiheadAttention's multi_head_attention_forward gained an
SDPA-branched implementation in torch 2.0. The branching caused the
tracer to specialize tgt_len as a Python int constant and bake it into
the output Reshape, so a model traced at T=40 errored with
'requested shape:{40,2,32}' at any other length. The historical
comment blaming espnet_positional_embedding.py was incorrect: the
failure reproduces with a bare nn.MultiheadAttention and zero PE code,
and survives even with the sinusoidal PE path which never touches
the espnet module.

Routed both non-RoPE paths through the in-house manual attention
(MultiheadSelfAttentionWithRoPE with rotary_embed=None) that was
already used on the RoPE path. It is fully dynamic-safe and produces
identical eager output (max-diff 7e-7) at T=40/80/160.

Checkpoint compatibility
Manual attention uses state_dict key 'in_proj.weight' whereas
nn.MultiheadAttention used 'in_proj_weight'. Same shape and same
Q/K/V-stacked-along-dim-0 semantics; utils.load_ckpt now renames the
old key on load, so legacy ckpts continue to work with strict=True.

Dependency cleanup

• Removed requirements-onnx.txt entirely (training and export share requirements.txt with PyTorch >= 2.4).
• Replaced onnxsim with onnxslim (>=0.1.93) via a thin utils.onnx_helper.simplify_onnx wrapper. onnxslim is easier to install across environments and has no native build chain.
• All torch.onnx.export calls now explicitly pass dynamo=False to stay on the TorchScript exporter that the existing graph-surgery in utils.onnx_helper was written against (silences the 2.9 default-switch warning).
• opset 15 -> 17. Verified with onnx.checker on all three model families.

[autumn-2-net]

@codex review

[yqzhishen]

Review (yqzhishen)

I've read the full diff and the existing review from HuanLinOTO. Here are my findings:

The existing blocking concern is already addressed

The reviewer flagged dynamo=False as incompatible with PyTorch 2.0–2.3 and 2.5–2.8. However, the actual code already implements the conditional gating:

TORCHSCRIPT_EXPORT_KWARGS: Dict[str, object] = (
    {'dynamo': False}
    if 'dynamo' in inspect.signature(torch.onnx.export).parameters
    else {}
)

inspect.signature probes the kwarg at import time. On versions without dynamo, TORCHSCRIPT_EXPORT_KWARGS is {} and **{} splats to nothing. On 2.4 and 2.9+ where the kwarg exists, dynamo=False is forwarded. No blocker — the code already handles every PyTorch version correctly.

Positive findings

Change	Verdict
WaveNet spec.squeeze(1) → spec[:, 0]	Correct. Avoids onnx::If rank mismatch.
Non-RoPE attention → MultiheadSelfAttentionWithRoPE	Correct. Unifies code paths, eliminates SDPA tgt_len baking. Numerical equivalence verified (max-diff ~7e-7).
in_proj_weight → in_proj.weight checkpoint rename	Correct. Minimal, covers all load_ckpt paths.
onnxsim → onnxslim wrapper	Clean. No native build chain.
opset 15 → 17	Reasonable. All used ops supported in 17.
Removed pre-surgery simplifier pass	Correct. Avoids topology collision with conditioner extraction.
TORCHSCRIPT_EXPORT_KWARGS conditional	Correct. Handles fragmented dynamo kwarg across all PyTorch versions.
requirements-onnx.txt removal	Correct. Single environment simplifies maintenance.

Minor notes (non-blocking)

• The assert check safety net from onnxsim.simplify is gone — onnxslim's slim() doesn't return a validation flag. Low risk.
• opset 17 was verified with onnx.checker; a full inference sanity check on all model families would be ideal for extra confidence.

Verdict

The PR is solid. All three bugs are correctly diagnosed and fixed. The reviewer's blocking concern is a false alarm. Approved.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 641f0bf5a9

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".