code: 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
  
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 3 for calling a global function with example usage: basic_pool(features, max_area_width, max_area_height, height, fn, name) and returns: pool_results,  area_heights,  area_widths