code: 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
  
file_path: tensor2tensor/tensor2tensor/layers/area_attention.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from six.moves import range
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
label_desc: Use image node_id 4 for calling a global function with example usage: _compute_sum_image(features, max_area_width, max_area_height, height, name) and returns: sum_image,  area_heights,  area_widths