test(stft): add return_complex=false and win_length < n_fft alignment tests

- write_stft_result_to_file: support kFloat/kDouble (real tensor from return_complex=false)
- Add ReturnComplexFalse test: validates [batch, freq, frames, 2] shape output
- Add WinLengthSmallerThanNFFT test: validates zero-pad window behavior

Author: youge325

test/ATen/ops/StftTest.cpp

@@ -1,7 +1,6 @@
 #include <ATen/ATen.h>
 #include <ATen/core/Tensor.h>
 #include <ATen/ops/ones.h>
-#include <ATen/ops/stft.h>
 #include <gtest/gtest.h>
 
 #include <cmath>
@@ -46,6 +45,16 @@ static void write_stft_result_to_file(FileManerger* file,
       std::complex<double>* data =
           reinterpret_cast<std::complex<double>*>(contig.data_ptr());
       *file << std::to_string(data[f * freq_bins].real()) << " ";
+    } else if (contig.scalar_type() == at::kFloat) {
+      // return_complex=false: shape [..., freq, frames, 2], last dim = [real,
+      // imag]
+      float* data = reinterpret_cast<float*>(contig.data_ptr());
+      int64_t stride = freq_bins * 2;
+      *file << std::to_string(data[f * stride]) << " ";
+    } else if (contig.scalar_type() == at::kDouble) {
+      double* data = reinterpret_cast<double*>(contig.data_ptr());
+      int64_t stride = freq_bins * 2;
+      *file << std::to_string(data[f * stride]) << " ";
     } else {
       *file << "unsupported_dtype ";
     }
@@ -214,5 +223,42 @@ TEST_F(StftTest, DifferentHopLength) {
   file.saveFile();
 }
 
+TEST_F(StftTest, ReturnComplexFalse) {
+  auto file_name = g_custom_param.get();
+  FileManerger file(file_name);
+  file.openAppend();
+  file << "ReturnComplexFalse ";
+  at::Tensor result = test_tensor.stft(
+      /*n_fft=*/8,
+      /*hop_length=*/4,
+      /*win_length=*/8,
+      /*window=*/::std::nullopt,
+      /*normalized=*/false,
+      /*onesided=*/true,
+      /*return_complex=*/false);
+  write_stft_result_to_file(&file, result);
+  file << "\n";
+  file.saveFile();
+}
+
+TEST_F(StftTest, WinLengthSmallerThanNFFT) {
+  auto file_name = g_custom_param.get();
+  FileManerger file(file_name);
+  file.openAppend();
+  file << "WinLengthSmallerThanNFFT ";
+  at::Tensor window = at::ones({4}, at::kFloat);
+  at::Tensor result = test_tensor.stft(
+      /*n_fft=*/8,
+      /*hop_length=*/4,
+      /*win_length=*/4,
+      /*window=*/window,
+      /*normalized=*/false,
+      /*onesided=*/true,
+      /*return_complex=*/true);
+  write_stft_result_to_file(&file, result);
+  file << "\n";
+  file.saveFile();
+}
+
 }  // namespace test
 }  // namespace at