self.window = torch.from_numpy(get_window(window, win_length, fftbins=True).astype(np.float32))？？？

[AutumnRoom]

How can a STFT be defined in the Burn like this one:

class TorchSTFT(nn.Module):
def init(self, filter_length=800, hop_length=200, win_length=800, window='hann'):
super().init()
self.filter_length = filter_length
self.hop_length = hop_length
self.win_length = win_length
self.window = torch.from_numpy(get_window(window, win_length, fftbins=True).astype(np.float32))

def transform(self, input_data):
    forward_transform = torch.stft(
        input_data,
        self.filter_length, self.hop_length, self.win_length, window=self.window.to(input_data.device),
        return_complex=True)
    return torch.abs(forward_transform), torch.angle(forward_transform)

def inverse(self, magnitude, phase):
    inverse_transform = torch.istft(
        magnitude * torch.exp(phase * 1j),
        self.filter_length, self.hop_length, self.win_length, window=self.window.to(magnitude.device))
    return inverse_transform.unsqueeze(-2)  # unsqueeze to stay consistent with conv_transpose1d implementation

def forward(self, input_data):
    self.magnitude, self.phase = self.transform(input_data)
    reconstruction = self.inverse(self.magnitude, self.phase)
    return reconstruction

[laggui]

Not familiar with this scipy function, but turns out there is also a pytorch equivalent.

You can implement it with Burn given the formula.

fn hann_window<B: Backend>(
    mut win_length: usize,
    periodic: bool,
    device: &B::Device,
) -> Tensor<B, 1> {
    // Tensor::<B, 1>::arang
    let n = Tensor::arange(0..win_length as i64, device).float();

    if periodic {
        win_length += 1;
    }

    // 0.5 - 0.5 * cos(2 * PI * n / (win_length - 1)))
    n.mul_scalar(2. * core::f32::consts::PI / (win_length - 1) as f32)
        .cos()
        .mul_scalar(-0.5)
        .add_scalar(0.5)
}