@parameterized.parameters((64, 16, 8, 16, np.float32, 0.002, 0.0005), (64, 16, 8, 16, np.float64, 1e-08, 1e-08), (64, 16, 16, 16, np.float32, 0.002, 0.0005), (64, 16, 16, 16, np.float64, 1e-08, 1e-08), (64, 16, 7, 16, np.float32, 0.002, 0.0005), (64, 16, 7, 16, np.float64, 1e-08, 1e-08), (64, 7, 4, 9, np.float32, 0.002, 0.0005), (64, 7, 4, 9, np.float64, 1e-08, 1e-08), (29, 5, 1, 10, np.float32, 0.002, 0.0005), (29, 5, 1, 10, np.float64, 1e-08, 1e-08))
@test.disable_with_predicate(pred=test.is_built_with_rocm, skip_message='On ROCm, this fails with mismatches at some locations (possibly due to peculiarities of rocFFT - investigate)')
def test_gradients_numerical(self, signal_length, frame_length, frame_step, fft_length, np_rtype, forward_tol, backward_tol):
    signal = ((np.random.rand(signal_length).astype(np_rtype) * 2) - 1)

    def forward(signal):
        return spectral_ops.stft(signal, frame_length, frame_step, fft_length, pad_end=False)
    ((f_jacob_t,), (f_jacob_n,)) = gradient_checker_v2.compute_gradient(forward, [signal])
    self.assertAllClose(f_jacob_t, f_jacob_n, rtol=forward_tol, atol=forward_tol)

    def backward(stft):
        return spectral_ops.inverse_stft(stft, frame_length, frame_step, fft_length)
    stft = forward(signal)
    ((b_jacob_t,), (b_jacob_n,)) = gradient_checker_v2.compute_gradient(backward, [stft])
    self.assertAllClose(b_jacob_t, b_jacob_n, rtol=backward_tol, atol=backward_tol)
@parameterized.parameters((64, 16, 8, 16, np.float32, 0.002, 0.0005), (64, 16, 8, 16, np.float64, 1e-08, 1e-08), (64, 16, 16, 16, np.float32, 0.002, 0.0005), (64, 16, 16, 16, np.float64, 1e-08, 1e-08), (64, 16, 7, 16, np.float32, 0.002, 0.0005), (64, 16, 7, 16, np.float64, 1e-08, 1e-08), (64, 7, 4, 9, np.float32, 0.002, 0.0005), (64, 7, 4, 9, np.float64, 1e-08, 1e-08), (29, 5, 1, 10, np.float32, 0.002, 0.0005), (29, 5, 1, 10, np.float64, 1e-08, 1e-08))
@test.disable_with_predicate(pred=test.is_built_with_rocm, skip_message='On ROCm, this fails with mismatches at some locations (possibly due to peculiarities of rocFFT - investigate)')
def test_gradients_numerical(self, signal_length, frame_length, frame_step, fft_length, np_rtype, forward_tol, backward_tol):
    signal = ((np.random.rand(signal_length).astype(np_rtype) * 2) - 1)

    def forward(signal):
        return spectral_ops.stft(signal, frame_length, frame_step, fft_length, pad_end=False)
    ((f_jacob_t,), (f_jacob_n,)) = gradient_checker_v2.compute_gradient(forward, [signal])
    self.assertAllClose(f_jacob_t, f_jacob_n, rtol=forward_tol, atol=forward_tol)

    def backward(stft):
        return spectral_ops.inverse_stft(stft, frame_length, frame_step, fft_length)
    stft = forward(signal)
    ((b_jacob_t,), (b_jacob_n,)) = gradient_checker_v2.compute_gradient(backward, [stft])
    self.assertAllClose(b_jacob_t, b_jacob_n, rtol=backward_tol, atol=backward_tol)