Fix STFT complex input frame offsets

onnxruntime/core/providers/cpu/signal/dft.cc

@@ -524,14 +524,15 @@ template <typename T, typename U>
 static Status short_time_fourier_transform(OpKernelContext* ctx, bool is_onesided, bool /*inverse*/) {
   // Attr("onesided"): default = 1
   // Input(0, "signal") type = T1
-  // Input(1, "frame_length") type = T2
+  // Input(1, "frame_step") type = T2
   // Input(2, "window") type = T1, optional
-  // Input(3, "frame_step") type = T2
+  // Input(3, "frame_length") type = T2
   // Output(0, "output") type = T1
 
   // Get signal
   const auto* signal = ctx->Input<Tensor>(0);
   const auto frame_step = signal::get_scalar_value_from_tensor<int64_t>(ctx->Input<Tensor>(1));
+  ORT_RETURN_IF_NOT(frame_step > 0, "frame_step must be greater than zero.");
   const auto* window = ctx->Input<Tensor>(2);
   const auto* frame_length_tensor = ctx->Input<Tensor>(3);
 
@@ -596,8 +597,11 @@ static Status short_time_fourier_transform(OpKernelContext* ctx, bool is_oneside
   // Run each dft of each batch as if it was a real-valued batch size 1 dft operation
   for (int64_t batch_idx = 0; batch_idx < batch_size; batch_idx++) {
     for (int64_t i = 0; i < n_dfts; i++) {
-      auto input_frame_begin =
-          signal_data + (batch_idx * signal_size * signal_components) + (i * frame_step * signal_components);
+      const auto frame_start = i * frame_step;
+      // Defensive check before creating a non-owning tensor view. n_dfts derivation should keep this in bounds.
+      ORT_RETURN_IF_NOT(frame_start <= signal_size - window_size, "STFT input frame is out of bounds.");
+      // signal_data is U*, so one increment advances one input sample, including both lanes for complex input.
+      auto input_frame_begin = signal_data + (batch_idx * signal_size) + frame_start;
 
       auto output_frame_begin = Y_data + (batch_idx * n_dfts * dft_output_size * output_components) +
                                 (i * dft_output_size * output_components);
@@ -619,9 +623,9 @@ static Status short_time_fourier_transform(OpKernelContext* ctx, bool is_oneside
 Status STFT::Compute(OpKernelContext* ctx) const {
   // Attr("onesided"): default = 1
   // Input(0, "signal") type = T1
-  // Input(1, "frame_length") type = T2
+  // Input(1, "frame_step") type = T2
   // Input(2, "window") type = T1, optional
-  // Input(3, "frame_step") type = T2
+  // Input(3, "frame_length") type = T2
   // Output(0, "output") type = T1
 
   // Get signal shape

onnxruntime/test/providers/cpu/signal/signal_ops_test.cc

@@ -233,6 +233,80 @@ TEST(SignalOpsTest, STFTFloat) {
   test.Run();
 }
 
+static void TestSTFTInvalidFrameStep(int64_t frame_step) {
+  OpTester test("STFT", kMinOpsetVersion);
+
+  vector<float> signal(64, 1);
+  test.AddInput<float>("signal", {1, 64, 1}, signal);
+  test.AddInput<int64_t>("frame_step", {}, {frame_step});
+  vector<float> window(16, 1);
+  test.AddInput<float>("window", {16}, window);
+  test.AddInput<int64_t>("frame_length", {}, {16});
+
+  vector<int64_t> output_shape = {1, 7, 9, 2};
+  vector<float> expected_output(1 * 7 * 9 * 2, 0.f);
+  test.AddOutput<float>("output", output_shape, expected_output);
+  test.Config(OpTester::ExpectResult::kExpectFailure, "frame_step must be greater than zero");
+  test.ConfigExcludeEps({kDmlExecutionProvider});
+  test.RunWithConfig();
+}
+
+TEST(SignalOpsTest, STFTFrameStepMustBePositive) {
+  TestSTFTInvalidFrameStep(0);
+  TestSTFTInvalidFrameStep(-1);
+}
+
+template <typename T>
+static void TestSTFTComplexInputBatched() {
+  OpTester test("STFT", kMinOpsetVersion);
+  test.AddAttribute<int64_t>("onesided", static_cast<int64_t>(false));
+
+  constexpr int64_t batch_size = 2;
+  constexpr int64_t signal_size = 128;
+  constexpr int64_t signal_components = 2;
+  constexpr int64_t frame_length = 32;
+  constexpr int64_t frame_step = 16;
+  constexpr int64_t n_dfts = 7;
+  constexpr int64_t dft_output_size = frame_length;
+  constexpr int64_t output_components = 2;
+
+  vector<T> signal(batch_size * signal_size * signal_components, static_cast<T>(0));
+  for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
+    const T signal_value = batch_idx == 0 ? static_cast<T>(1) : static_cast<T>(99);
+    for (int64_t sample_idx = 0; sample_idx < signal_size; ++sample_idx) {
+      signal[(batch_idx * signal_size + sample_idx) * signal_components] = signal_value;
+    }
+  }
+
+  test.AddInput<T>("signal", {batch_size, signal_size, signal_components}, signal);
+  test.AddInput<int64_t>("frame_step", {}, {frame_step});
+  test.AddOptionalInputEdge<T>();
+  test.AddInput<int64_t>("frame_length", {}, {frame_length});
+
+  vector<int64_t> output_shape = {batch_size, n_dfts, dft_output_size, output_components};
+  vector<T> expected_output(batch_size * n_dfts * dft_output_size * output_components, static_cast<T>(0));
+  for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
+    const T dc_value = (batch_idx == 0 ? static_cast<T>(1) : static_cast<T>(99)) * static_cast<T>(frame_length);
+    for (int64_t frame_idx = 0; frame_idx < n_dfts; ++frame_idx) {
+      expected_output[((batch_idx * n_dfts + frame_idx) * dft_output_size) * output_components] = dc_value;
+    }
+  }
+
+  test.AddOutput<T>("output", output_shape, expected_output);
+  test.SetOutputAbsErr("output", 0.001f);
+  // DML does not consistently match these CPU STFT validation/regression paths in Windows GPU CI.
+  test.ConfigExcludeEps({kDmlExecutionProvider});
+  test.RunWithConfig();
+}
+
+TEST(SignalOpsTest, STFTFloatComplexInputBatched) {
+  TestSTFTComplexInputBatched<float>();
+}
+
+TEST(SignalOpsTest, STFTDoubleComplexInputBatched) {
+  TestSTFTComplexInputBatched<double>();
+}
+
 TEST(SignalOpsTest, HannWindowFloat) {
   OpTester test("HannWindow", kMinOpsetVersion);