Confused about tf.keras.layers.Flatten

The following example

import numpy as np import tensorflow as tf model = tf.keras.Sequential() model.add(tf.keras.layers.Dense(2, activation='relu', input_shape=(2,2,))) model.add(tf.keras.layers.Flatten()) model.add(tf.keras.layers.Dense(1)) tx = np.random.rand(2,2) res = model(tx) print(res) 

Gives error

ValueError: Input 0 of layer dense_1 is incompatible with the layer: expected axis -1 of input shape to have value 4 but received input with shape (2, 2) 

But if I comment out the line with Flatten layer, then everything works fine

What is wrong with this code and how do I properly flatten output layer?