Neural Network to learn the Discrete Fourier Transform

dft_nn.py

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf

batch_size = 15000
fft_size = 128

xs = np.random.randn(batch_size, fft_size) + np.random.randn(batch_size, fft_size) * 1j
ys = np.fft.fft(xs, axis=1)

(Xs, Ys) = (np.stack([xs.real, xs.imag], axis=2), np.hstack([ys.real, ys.imag]))
print(Xs.shape, Ys.shape)

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(fft_size,2)),
    tf.keras.layers.Dense(fft_size*2, use_bias=False)
])

optimizer = tf.keras.optimizers.Adam(amsgrad=True)
model.compile(loss='mean_squared_error', optimizer=optimizer)
history = model.fit(Xs, Ys, epochs=100)