Implement aten.stft

[moatom]

Fixed pytorch/pytorch#147052

$ python -m pytest tests/function_libs/torch_lib/ops_test.py -k ops_aten_stft
====================================================================================================================================================================================================== test session starts ======================================================================================================================================================================================================
platform linux -- Python 3.13.1, pytest-8.4.1, pluggy-1.6.0
Using --randomly-seed=371864411
rootdir: /home/moatom/github/onnxscript
configfile: pyproject.toml
plugins: randomly-3.16.0, xdist-3.8.0, subtests-0.14.2, cov-6.2.1, hypothesis-6.138.2
collected 2158 items / 2154 deselected / 4 selected                                                                                                                                                                                                                                                                                                                                                                             

tests/function_libs/torch_lib/ops_test.py s..x                                                                                                                                                                                                                                                                                                                                                                     [100%]

======================================================================================================================================================================================================= warnings summary ========================================================================================================================================================================================================
onnxscript/converter.py:457: 429 warnings
tests/function_libs/torch_lib/ops_test.py: 15 warnings
  /home/moatom/github/onnxscript/onnxscript/converter.py:457: DeprecationWarning: Expression.__init__ got an unexpected keyword argument 'lineno'. Support for arbitrary keyword arguments is deprecated and will be removed in Python 3.15.
    expr = ast.Expression(expr, lineno=expr.lineno, col_offset=expr.col_offset)

onnxscript/converter.py:457: 429 warnings
tests/function_libs/torch_lib/ops_test.py: 15 warnings
  /home/moatom/github/onnxscript/onnxscript/converter.py:457: DeprecationWarning: Expression.__init__ got an unexpected keyword argument 'col_offset'. Support for arbitrary keyword arguments is deprecated and will be removed in Python 3.15.
    expr = ast.Expression(expr, lineno=expr.lineno, col_offset=expr.col_offset)

tests/function_libs/torch_lib/ops_test.py::TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__ops_aten_stft_cpu_float32
tests/function_libs/torch_lib/ops_test.py::TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__ops_aten_stft_cpu_float32
tests/function_libs/torch_lib/ops_test.py::TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__ops_aten_stft_cpu_float32
  /home/moatom/github/onnxscript/tests/function_libs/torch_lib/ops_test_common.py:329: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. To learn more, see the migration guide https://numpy.org/devdocs/numpy_2_0_migration_guide.html#adapting-to-changes-in-the-copy-keyword
    value = np.array(value.cpu())

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
==================================================================================================================================================================================================== short test summary info ====================================================================================================================================================================================================
SKIPPED [1] tests/function_libs/torch_lib/ops_test.py:101: Traced functions does not have a function proto
=================================================================================================================================================================== 2 passed, 1 skipped, 2154 deselected, 1 xfailed, 891 warnings, 7 subtests passed in 4.42s ===================================================================================================================================================================

[moatom]

@microsoft-github-policy-service agree

[moatom]

pytorch/pytorch#147052 (comment)

It would be nice to decay to onnx.STFT node in the graph if options of a given call allow this...

I think this PR’s implementation already uses onnx.STFT as its core: https://github.com/microsoft/onnxscript/pull/2645/files#diff-773a426eaaf0eb2a1ed6039f2cd50613d2db7d61c239453545413499b2a15776R8194

result = op.STFT(signal, frame_step_const, window, frame_length_const, onesided=onesided)

So I don’t think an decay to onnx.STFT is necessary in this case. 🤔

[moatom]

@justinchuby

Hi!

pytorch/pytorch#147052 (comment)

I think it can be simplified quite a bit.

Could you elaborate a bit on this comment at your convenience? (If you think this simplification is necessary.)

[Copilot]

Pull Request Overview

This PR implements the aten::stft (Short-Time Fourier Transform) operator to resolve issue #147052. The implementation includes handling for various optional parameters like hop_length, win_length, window, normalized, onesided, and return_complex.

Key changes:

• Added STFT operator implementation with helper functions for batch dimension handling, window centering, and FFT normalization
• Registered the operator in test data with appropriate tolerance settings and xfail for float16 dtype

Reviewed Changes

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

File	Description
onnxscript/function_libs/torch_lib/ops/core.py	Implements aten_stft and five helper functions for STFT processing
tests/function_libs/torch_lib/ops_test_data.py	Registers the new operator in test suite with tolerance and xfail configuration

[codecov]

Codecov Report

❌ Patch coverage is 65.45455% with 19 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.43%. Comparing base (5be9d3b) to head (c0fd969).

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	65.45%	15 Missing and 4 partials ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2645   +/-   ##
=======================================
  Coverage   70.43%   70.43%           
=======================================
  Files         224      224           
  Lines       26778    26833   +55     
  Branches     2680     2689    +9     
=======================================
+ Hits        18862    18901   +39     
- Misses       6985     6997   +12     
- Partials      931      935    +4     

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

[titaiwangms]

Some minor comments and this should be the final pass.
cc @justinchuby I remember you said something stft implementation was wrong. This is heavily following the original one. Would you care share about your comments?

[justinchuby]

Thanks - I need to read this more closely to remind myself what was going on then

[justinchuby]

Thanks. I don't have a better idea

[titaiwangms]

@moatom Could you do a final pass, and make sure lint is passing.

[moatom]

@titaiwangms I resolved all of them. Could you please confirm when you have a moment?

[justinchuby]

If possible, could you help add a few tests to https://github.com/microsoft/onnxscript/blob/main/tests/function_libs/torch_lib/e2e_ops_tests.py ? Thanks a lot!

[moatom]

@justinchuby

If possible, could you help add a few tests to https://github.com/microsoft/onnxscript/blob/main/tests/function_libs/torch_lib/e2e_ops_tests.py ?

DONE