zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 17:31:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Formatted

Browse Source
This commit is contained in:
paramthakkar123 2025-04-17 00:24:49 +05:30
parent 9e8befbe8d
commit f50cce83a5
1 changed files with 30 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
python/mlx/stft.py
View File

@ -1,6 +1,7 @@
import mlx.core as mx import mlx.core as mx
import numpy as np import numpy as np
def stft( def stft(
x: mx.array, x: mx.array,
n_fft: int = 2048, n_fft: int = 2048,
@ -34,10 +35,9 @@ def stft(
n_frames = 1 + (x.shape[0] - n_fft) // hop_length n_frames = 1 + (x.shape[0] - n_fft) // hop_length
frames = mx.stack([ frames = mx.stack(
x[i * hop_length : i * hop_length + n_fft] * window [x[i * hop_length : i * hop_length + n_fft] * window for i in range(n_frames)]
for i in range(n_frames) )
])
stft = mx.fft.fft(frames, n=n_fft, axis=-1) stft = mx.fft.fft(frames, n=n_fft, axis=-1)
@ -52,6 +52,7 @@ def stft(
return stft return stft
def istft( def istft(
stft_matrix: mx.array, stft_matrix: mx.array,
hop_length: int = None, hop_length: int = None,