code: def _pool_one_shape(
      features_2d,
      area_width,
      area_height,
      batch_size,
      width,
      height,
      depth,
      fn=tf.reduce_max,
      name=None,
  ):
      """Pools for an area in features_2d.
  
      Args:
        features_2d: a Tensor in a shape of [batch_size, height, width, depth].
        area_width: the max width allowed for an area.
        area_height: the max height allowed for an area.
        batch_size: the batch size.
        width: the width of the memory.
        height: the height of the memory.
        depth: the depth of the features.
        fn: the TF function for the pooling.
        name: the op name.
      Returns:
        pool_tensor: A Tensor of shape [batch_size, num_areas, depth]
      """
      with tf.name_scope(name, default_name="pool_one_shape"):
          images = []
          for y_shift in range(area_height):
              image_height = tf.maximum(
                  height - area_height + 1 + y_shift, 0
              )
              for x_shift in range(area_width):
                  image_width = tf.maximum(
                      width - area_width + 1 + x_shift, 0
                  )
                  area = features_2d[
                      :, y_shift:image_height, x_shift:image_width, :
                  ]
                  flatten_area = tf.reshape(
                      area, [batch_size, -1, depth, 1]
                  )
                  images.append(flatten_area)
          image_tensor = tf.concat(images, axis=3)
          max_tensor = fn(image_tensor, axis=3)
      return max_tensor
  
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 2 for calling a global function with example usage: _pool_one_shape(features_2d, area_width, area_height, batch_size, width, height, depth, fn, name) and returns: max_tensor