ITADN

GRU translation missing linear_before_reset=1 (produces incorrect results for PyTorch GRU)

#2852Opendaniel-om-weber 创建于 2026-03-14
D
daniel-om-webercommented
## Summary The `aten_gru` translation in `onnxscript/function_libs/torch_lib/ops/core.py` (added in #2674) does not set `linear_before_reset=1` on the ONNX `GRU` op. This causes numerically incorrect results because PyTorch's `nn.GRU` uses the `linear_before_reset=1` variant. ## Details PyTorch GRU computes the new gate as: ``` n_t = tanh(W_in @ x_t + b_in + r_t * (W_hn @ h_{t-1} + b_hn)) ``` This matches the ONNX GRU spec with `linear_before_reset=1`. But the default `linear_before_reset=0` applies the reset gate *before* the linear transformation — a different equation. The two `op.GRU` calls (~lines 4352 and 4362) need `linear_before_reset=1` added. ## Reproduction ```python import torch, numpy as np m = torch.nn.GRU(1, 32, batch_first=True) m.eval() inp = torch.randn(1, 10, 1) with torch.no_grad(): pt_out, _ = m(inp) torch.onnx.export(m, (inp,), f="gru.onnx") import onnxruntime as ort sess = ort.InferenceSession("gru.onnx") onnx_out = sess.run(None, {sess.get_inputs()[0].name: inp.numpy()})[0] print("Max abs diff:", np.abs(pt_out.numpy() - onnx_out).max()) # Expected: ~1e-7 (float32 precision) # Actual: ~0.1 (incorrect GRU equation) ``` ## Environment - torch 2.10.0 - onnxscript 0.6.2 - onnxruntime 1.22.0 ## References - [ONNX GRU spec (`linear_before_reset`)](https://onnx.ai/onnx/operators/onnx__GRU.html) - [PyTorch GRU docs](https://pytorch.org/docs/stable/generated/torch.nn.GRU.html) - PR #2674 (introduced the `aten_gru` translation)
2 条评论