code: def batch_flatten(x):
      
      shape = x.get_shape().as_list()[1:]
      if None not in shape:
          return tf.reshape(x, [-1, int(np.prod(shape))])
      return tf.reshape(x, tf.stack([tf.shape(x)[0], -1]))
docs: Flatten the tensor except the first dimension.