code: def flatten(input_layer, preserve_batch=True):
    
    if preserve_batch:
      return reshape(input_layer, [DIM_SAME, -1])
    else:
      return reshape(input_layer, [-1])
docs: Flattens this.
  
    If preserve_batch is True, the result is rank 2 and the first dim (batch) is
    unchanged. Otherwise the result is rank 1.
  
    Args:
      input_layer: The Pretty Tensor object, supplied.
      preserve_batch: If True (the default), then preserve the first dimension.
    Returns:
      A LayerWrapper with the flattened tensor.