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