Computes area sums for features.

  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.
    name: the namescope.
  Returns:
    sum_image: 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 _compute_sum_image(features, max_area_width, max_area_height=1, height=1,
                       name=None):
  """Computes area sums for features.

  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.
    name: the namescope.
  Returns:
    sum_image: 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="compute_sum_image"):
    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])
    width_cum = tf.cumsum(features_2d, axis=-2, name="compute_integral_h")
    integral_image = tf.cumsum(width_cum, axis=-3, name="compute_integral_v")
    padded_image = tf.pad(
        integral_image, [[0, 0], [1, 0], [1, 0], [0, 0]], constant_values=0)
    height_list = []
    width_list = []
    dst_images = []
    src_images_diag = []
    src_images_h = []
    src_images_v = []
    size_tensor = tf.ones_like(padded_image[:, :, :, 0],
                               dtype=tf.int32)
    for area_height in range(max_area_height):
      for area_width in range(max_area_width):
        dst_images.append(
            tf.reshape(
                padded_image[:, area_height + 1:, area_width + 1:, :],
                [batch_size, -1, depth]))
        src_images_diag.append(
            tf.reshape(
                padded_image[:, :-area_height - 1, :-area_width - 1, :],
                [batch_size, -1, depth]))
        src_images_h.append(
            tf.reshape(
                padded_image[:, area_height + 1:, :-area_width - 1, :],
                [batch_size, -1, depth]))
        src_images_v.append(
            tf.reshape(
                padded_image[:, :-area_height - 1, area_width + 1:, :],
                [batch_size, -1, depth]))
        height_list.append(
            tf.reshape(
                size_tensor[:, area_height + 1:, area_width + 1:] *\
                (area_height + 1), [batch_size, -1]))
        width_list.append(
            tf.reshape(
                size_tensor[:, area_height + 1:, area_width + 1:] *\
                (area_width + 1), [batch_size, -1]))
    sum_image = tf.subtract(
        tf.concat(dst_images, axis=1) + tf.concat(src_images_diag, axis=1),
        tf.concat(src_images_v, axis=1) + tf.concat(src_images_h, 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 sum_image, area_heights, area_widths