# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import torch
from librosa import stft
import numpy as np
import matplotlib.pyplot as plt

# t = 2*np.pi*np.linspace(0,1,5000)
# x = np.sin(50*t)
x = np.random.randn(10,10000).astype(np.float32)
lib_output = stft(x[0], n_fft=2048, hop_length=512)
window = torch.hann_window(2048)

x = torch.nn.Parameter(x)

torch_output = torch.stft(x, n_fft=2048, hop_length=512, window=window)

plt.imshow(torch.sqrt(torch_output[0].pow(2).sum(-1)), aspect='auto')

plt.imshow(abs(lib_output), aspect='auto')

np.allclose(torch.sqrt(torch_output[0].pow(2).sum(-1)), abs(lib_output), atol=1e-5, rtol=1e-6)



StarcoderdataJupyterScriptsDedupFiltered