Let's load the MNIST dataset. Since we are reconstructing the given input, we don't need the labels. So, we just load x_train for training and x_test for testing:

(x_train, _), (x_test, _) = mnist.load_data()

Normalize the data by dividing by the max pixel value, which is 255:

x_train = x_train.astype('float32') / 255
x_test = x_test.astype('float32') / 255

Print the shape of our dataset:

print(x_train.shape, x_test.shape)

((60000, 28, 28), (10000, 28, 28))

Reshape the images as a 2D array:

x_train = x_train.reshape((len(x_train), np.prod(x_train.shape[1:])))
x_test = x_test.reshape((len(x_test), np.prod(x_test.shape[1:])))

Now, the shape of the data would become as follows:

print(x_train.shape, x_test.shape)

((60000, 784), (10000, 784))