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()