Output shape of a flatten layer.

def _flatten_output_shape(input_shape, num_axis_to_keep=1):
  """Output shape of a flatten layer."""
  if num_axis_to_keep >= len(input_shape):
    raise ValueError(
        "num_axis_to_keep[%d] should be less than input's rank[%d]" %
        (num_axis_to_keep, len(input_shape)))
  return tuple(input_shape[:num_axis_to_keep]) + (
      reduce(op.mul, input_shape[num_axis_to_keep:], 1),)