code: def _flatten_output_shape(input_shape, num_axis_to_keep=1):
    
    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),)
docs: Output shape of a flatten layer.