Pools for each area based on a given pooling function (fn).

  Args:
    features: a Tensor in a shape of [batch_size, height * width, depth].
    max_area_width: the max width allowed for an area.
    max_area_height: the max height allowed for an area.
    height: the height of the image.
    fn: the TF function for the pooling.
    name: the namescope.
  Returns:
    pool_results: A Tensor of shape [batch_size, num_areas, depth]
    area_heights: A Tensor of shape [batch_size, num_areas, 1]
    area_widths: A Tensor of shape [batch_size, num_areas, 1]

def basic_pool(features, max_area_width, max_area_height=1, height=1,
               fn=tf.reduce_max, name=None):
  """Pools for each area based on a given pooling function (fn).

  Args:
    features: a Tensor in a shape of [batch_size, height * width, depth].
    max_area_width: the max width allowed for an area.
    max_area_height: the max height allowed for an area.
    height: the height of the image.
    fn: the TF function for the pooling.
    name: the namescope.
  Returns:
    pool_results: A Tensor of shape [batch_size, num_areas, depth]
    area_heights: A Tensor of shape [batch_size, num_areas, 1]
    area_widths: A Tensor of shape [batch_size, num_areas, 1]
  """
  with tf.name_scope(name, default_name="basic_pool"):
    feature_shape = common_layers.shape_list(features)
    batch_size = feature_shape[0]
    length = feature_shape[-2]
    depth = feature_shape[-1]
    width = length // height
    features_2d = tf.reshape(features, [batch_size, height, width, depth])
    height_list = []
    width_list = []
    pool_list = []
    size_tensor = tf.ones_like(features_2d[:, :, :, 0], dtype=tf.int32)
    for area_height in range(max_area_height):
      for area_width in range(max_area_width):
        pool_tensor = _pool_one_shape(features_2d,
                                      area_width=area_width + 1,
                                      area_height=area_height + 1,
                                      batch_size=batch_size,
                                      width=width,
                                      height=height,
                                      depth=depth,
                                      fn=fn)
        pool_list.append(
            tf.reshape(pool_tensor, [batch_size, -1, depth]))
        height_list.append(
            tf.reshape(
                size_tensor[:, area_height:, area_width:] *\
                (area_height + 1), [batch_size, -1]))
        width_list.append(
            tf.reshape(
                size_tensor[:, area_height:, area_width:] *\
                (area_width + 1), [batch_size, -1]))
    pool_results = tf.concat(pool_list, axis=1)
    area_heights = tf.expand_dims(tf.concat(height_list, axis=1), 2)
    area_widths = tf.expand_dims(tf.concat(width_list, axis=1), 2)
  return pool_results, area_heights, area_widths