def test_fft_forward_backward(seed, ctx, func_name, batch_dims, signal_ndim, dims, normalized): if func_name == "IFFT": pytest.skip("Not implemented in CPU.") from nbla_test_utils import function_tester, convert_to_float2_array, convert_to_complex_array rng = np.random.RandomState(seed) shape = batch_dims + dims x_data_complex = rng.rand(*shape) + 1j * rng.rand(*shape) x_data = convert_to_float2_array(x_data_complex) inputs = [x_data] func_args = [signal_ndim, normalized] function_tester(rng, F.ifft, ref_ifft, inputs, func_args=func_args, atol_f=1e-4, atol_b=1e-4, backward=[True], ctx=ctx, func_name=func_name, ref_grad=ref_grad_ifft, disable_half_test=True)
def test_fft_double_backward(seed, ctx, func_name, batch_dims, signal_ndim, dims, normalized): if func_name == "IFFTCuda" and sys.platform == 'win32': from nnabla_ext import cuda if cuda._version.__cuda_version__ == '11.4': pytest.skip("Skip win32+CUDA114 tests") if func_name == "IFFT": pytest.skip("Not implemented in CPU.") from nbla_test_utils import backward_function_tester, convert_to_float2_array, convert_to_complex_array rng = np.random.RandomState(seed) shape = batch_dims + dims x_data_complex = rng.rand(*shape) + 1j * rng.rand(*shape) x_data = convert_to_float2_array(x_data_complex) inputs = [x_data] func_args = [signal_ndim, normalized] backward_function_tester(rng, F.ifft, inputs, func_args=func_args, atol_f=1e-4, atol_accum=5e-2, backward=[True], ctx=ctx)
def ref_fft(x, signal_ndim, normalized): from nbla_test_utils import convert_to_float2_array, convert_to_complex_array x_data_complex = convert_to_complex_array(x) batch_dims = x_data_complex.shape[0:len( x_data_complex.shape) - signal_ndim] ref_data_complex = np.fft.fftn(x_data_complex, axes=np.arange(signal_ndim) + len(batch_dims), norm="ortho" if normalized else None) ref_data_float2 = convert_to_float2_array( ref_data_complex).astype(np.float32) return ref_data_float2
def ref_grad_fft(x, dy, signal_ndim, normalized): from nbla_test_utils import convert_to_float2_array, convert_to_complex_array dy_complex = convert_to_complex_array(dy) batch_dims = dy_complex.shape[0:len(dy_complex.shape) - signal_ndim] ref_grad_complex = np.fft.ifftn(dy_complex, axes=np.arange(signal_ndim) + len(batch_dims), norm="ortho" if normalized else None) if not normalized: scale = np.prod(ref_grad_complex.shape[len(batch_dims):]) if len( batch_dims) > 0 else np.prod(ref_grad_complex.shape) ref_grad_complex *= scale ref_grad = convert_to_float2_array(ref_grad_complex).astype(np.float32) return ref_grad.flatten()
def test_fft_double_backward(seed, ctx, func_name, batch_dims, signal_ndim, dims, normalized): if func_name == "FFT": pytest.skip("Not implemented in CPU.") from nbla_test_utils import backward_function_tester, convert_to_float2_array, convert_to_complex_array rng = np.random.RandomState(seed) shape = batch_dims + dims x_data_complex = rng.rand(*shape) + 1j * rng.rand(*shape) x_data = convert_to_float2_array(x_data_complex) inputs = [x_data] func_args = [signal_ndim, normalized] backward_function_tester(rng, F.fft, inputs, func_args=func_args, atol_f=1e-3, atol_accum=8e-2, backward=[True], ctx=ctx)