title: Implement a flatten layer in Tensorflow
body: <p>I am trying to implement a flatten layer using TensorFlow 2.2.0. I am following the instructions from Geron's book (2nd ed).
  As for a flatten layer, I first try to get the batch input shape and compute the new shape.
  But I have met this problem with tensor dimensions: <code>TypeError: Dimension value must be integer or None or have an __index__ method</code></p>
  <pre><code>import tensorflow as tf
  from tensorflow import keras
  (X_train, y_train), (X_test, y_test) = keras.datasets.fashion_mnist.load_data()
  input_shape = X_train.shape[1:]
  assert input_shape == (28, 28)
  
  class MyFlatten(keras.layers.Layer):
      def __init__(self, **kwargs):
          super().__init__(**kwargs)
  
      def build(self, batch_input_shape):
          super().build(batch_input_shape) 
  
      def call(self, X):
          X_shape = tf.shape(X)
          batch_size = X_shape[0]
          new_shape = tf.TensorShape([batch_size, X_shape[1]*X_shape[2]])
          return tf.reshape(X, new_shape)
  
      def get_config(self):
          base_config = super().get_config()
          return {**base_config}
  
  ## works fine on this example
  MyFlatten()(X_train[:10])
  
  ## fail when building a model
  input_ = keras.layers.Input(shape=[28, 28])
  fltten_ = MyFlatten()(input_)
  hidden1 = keras.layers.Dense(300, activation=&quot;relu&quot;)(fltten_)
  hidden2 = keras.layers.Dense(100, activation=&quot;relu&quot;)(hidden1)
  output = keras.layers.Dense(10, activation=&quot;softmax&quot;)(hidden2)
  model = keras.models.Model(inputs=[input_], outputs=[output])
  model.summary()
  </code></pre>