Fix FFT export truncation by preserving dft_length in rFFT lowering

[Copilot]

The FFT lowering path was exporting real-to-complex transforms without carrying the active signal length into the ONNX DFT node. In some conversions, that allowed dft_length to default to 1, silently truncating the spectrum and producing incorrect output shapes for torch.fft.rfft.

• Summary

  • Port the upstream PyTorch fix into ONNXScript’s FFT lowering.
  • Preserve the true transformed axis length during aten::_fft_r2c export so the generated ONNX graph matches PyTorch FFT semantics.

• Changes

  • In onnxscript/function_libs/torch_lib/ops/fft.py, compute the current axis length from the tensor shape and pass it explicitly as dft_length when calling ONNX DFT.
  • Keep the existing normalization and one-sided handling intact while removing the implicit truncation bug.
  • Add a regression in tests/function_libs/torch_lib/e2e_ops_tests.py covering torch.fft.rfft(x, n=512, dim=-1) and asserting the exported output shape remains [4, 257].

x = torch.randn(4, 512)
y = torch.fft.rfft(x, n=512, dim=-1)
# expected shape: [4, 257]

• Fixes Port fix from pytorch for fft #2990
• Fix ONNX export via Dynamo sets dft_length = 1 in DFT, breaking shape-inference for torch.fft.rfft pytorch/pytorch#155997

[Copilot]

Pull request overview

This PR fixes an ONNX export correctness issue in the Torch FFT lowering path by explicitly preserving and forwarding the effective transform length (dft_length) into ONNX DFT during real-to-complex (rFFT) lowering, preventing silent spectrum truncation and incorrect exported output shapes.

Changes:

• Compute the active transform-axis length from the input tensor shape during aten::_fft_r2c lowering and pass it explicitly as dft_length to ONNX DFT.
• Feed the same computed dft_length into the normalization helper to keep scaling consistent with the transformed signal length.
• Add an end-to-end regression test asserting the exported output shape for torch.fft.rfft(x, n=512, dim=-1) is [4, 257].

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxscript/function_libs/torch_lib/ops/fft.py	Preserves the effective signal length by computing and passing dft_length into ONNX DFT during rFFT lowering.
tests/function_libs/torch_lib/e2e_ops_tests.py	Adds a regression test to ensure exported rFFT output shapes are not truncated.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 72.63%. Comparing base (ab2b5f4) to head (c5f15bc).
✅ All tests successful. No failed tests found.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2991   +/-   ##
=======================================
  Coverage   72.63%   72.63%           
=======================================
  Files         265      265           
  Lines       32208    32209    +1     
  Branches     3042     3042           
=======================================
+ Hits        23395    23396    +1     
  Misses       7779     7779           
  Partials     1034     1034           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.