import mindspore.numpy as np
def _fft_helper(x, win, detrend_func,
                nperseg: int, noverlap: int, nfft, sides: str):
  """Calculate windowed FFT in the same way the original SciPy does.
  """
  # if x.dtype.kind == 'i':
  #   x = x.astype(win.dtype)

  *batch_shape, signal_length = x.shape
  # Created strided array of data segments
  if nperseg == 1 and noverlap == 0:
    result = x[..., np.newaxis]
  else:
    step = nperseg - noverlap
    batch_shape = list(batch_shape)
    x = x.reshape((math.prod(batch_shape), signal_length, 1))
    print(f"nperseg {nperseg} step {step}")
    result = conv_general_dilated_patches_torch(x, nperseg, step)
    result = result.reshape(*batch_shape, *result.shape[-2:])

  # Detrend each data segment individually
  result = detrend_func(result)

  # Apply window by multiplication
  result = win.reshape((1,) * len(batch_shape) + (1, nperseg)) * result
  result = result.astype(ms.float32)
  # Perform the fft on last axis. Zero-pads automatically
  if sides == 'twosided':
    return np.fft.fft(result, n=nfft)
  else:
    return np.fft.rfft(result, n=nfft)

#简单测试一下求梯度
audio_data = ms.Tensor(audio_data,dtype=ms.float32)
def stft_wrap(audio_data):
    #stft主要运算部分为 fft_helper函数
    _,_,matrix = stft(audio_data, nfft=2048,noverlap=2048-512,nperseg=2048,boundary=None)
    matrix = (matrix.real()**2+matrix.imag()**2)**0.5
    return np.mean(matrix-1.)
grad_fn = ms.grad(stft_wrap, grad_position=(0),has_aux=False)
res = grad_fn(audio_data)

RuntimeError: aclnnAddGetWorkspaceSize call failed, please check!

----------------------------------------------------
- Ascend Error Message:
----------------------------------------------------
EZ1001: 2025-01-22-11:53:28.145.098 self not implemented for DT_COMPLEX64, should be in dtype support list [DT_FLOAT,DT_INT32,DT_INT64,DT_FLOAT16,DT_INT16,DT_INT8,DT_UINT8,DT_DOUBLE,DT_BOOL,DT_BFLOAT16,].[THREAD:12113]

(Please search "CANN Common Error Analysis" at https://www.mindspore.cn for error code description)

----------------------------------------------------
- C++ Call Stack: (For framework developers)
----------------------------------------------------
mindspore/ops/kernel/ascend/pyboost/customize/add.cc:44 operator()