code: def compute_area_features(
      features,
      max_area_width,
      max_area_height=1,
      height=1,
      epsilon=1e-6,
  ):
      """Computes features for each area.
  
      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.
        epsilon: the epsilon added to the variance for computing standard deviation.
      Returns:
        area_mean: A Tensor of shape [batch_size, num_areas, depth]
        area_std: A Tensor of shape [batch_size, num_areas, depth]
        area_sum: 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("compute_area_features"):
          tf.logging.info(
              "area_attention compute_area_features: %d x %d",
              max_area_height,
              max_area_width,
          )
          area_sum, area_heights, area_widths = _compute_sum_image(
              features,
              max_area_width=max_area_width,
              max_area_height=max_area_height,
              height=height,
          )
          area_squared_sum, _, _ = _compute_sum_image(
              tf.pow(features, 2),
              max_area_width=max_area_width,
              max_area_height=max_area_height,
              height=height,
          )
          sizes = tf.multiply(area_heights, area_widths)
          float_area_sizes = tf.to_float(sizes)
          area_mean = tf.div(area_sum, float_area_sizes)
          s2_n = tf.div(area_squared_sum, float_area_sizes)
          area_variance = tf.subtract(s2_n, tf.pow(area_mean, 2))
          area_std = tf.sqrt(tf.abs(area_variance) + epsilon)
          return (
              area_mean,
              area_std,
              area_sum,
              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 5 for calling a global function with example usage: compute_area_features(features, max_area_width, max_area_height, height, epsilon) and returns: area_mean,  area_std,  area_sum,  area_heights,  area_widths