code: def compute_area_key(
      features,
      max_area_width,
      max_area_height=1,
      height=1,
      mode="mean",
      training=True,
      name=None,
  ):
      """Computes the key 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.
        mode: whether to combine different area features or only use
            the vector mean of each area, which can be "mean", "concat", "sum",
            "sample_concat", and "sample_sum".
        training: indicating if it is in the training mode.
        name: the name for setting the variable scope.
      Returns:
        area_key: a Tensor in the shape of [batch_size, num_areas, depth]
      """
  
      tf.logging.info("area_attention mode=%s", mode)
      (
          area_mean,
          area_std,
          _,
          area_heights,
          area_widths,
      ) = compute_area_features(
          features,
          max_area_width=max_area_width,
          max_area_height=max_area_height,
          height=height,
      )
      if mode == "mean":
          return area_mean
      elif mode == "max":
          area_max, _, _ = basic_pool(
              features,
              max_area_width=max_area_width,
              max_area_height=max_area_height,
              height=height,
          )
          return area_max
      elif mode == "sample":
          if training:
              area_mean += area_std * tf.random_normal(
                  tf.shape(area_std)
              )
          return area_mean
      with tf.variable_scope(
          name,
          default_name="combine_area_features",
          values=[area_mean, area_std, area_heights, area_widths],
      ):
          depth = common_layers.shape_list(area_mean)[-1]
          height_embed = tf.nn.embedding_lookup(
              params=tf.get_variable(
                  "area_height_emb", [max_area_height, depth // 2]
              ),
              ids=area_heights[:, :, 0] - 1,
          )
          width_embed = tf.nn.embedding_lookup(
              params=tf.get_variable(
                  "area_width_emb", [max_area_width, depth // 2]
              ),
              ids=area_widths[:, :, 0] - 1,
          )
          size_embed = tf.concat([height_embed, width_embed], -1)
          if mode == "concat":
              feature_concat = tf.concat(
                  [area_mean, area_std, size_embed], -1
              )
          elif mode == "max_concat":
              area_max, _, _ = basic_pool(
                  features,
                  max_area_width=max_area_width,
                  max_area_height=max_area_height,
                  height=height,
              )
              feature_concat = tf.concat([area_max, size_embed], -1)
          elif mode == "sum":
              feature_concat = size_embed + area_mean + area_std
          elif mode == "sample_concat":
              if training:
                  area_mean += area_std * tf.random_normal(
                      tf.shape(area_std)
                  )
              feature_concat = tf.concat([area_mean, size_embed], -1)
          elif mode == "sample_sum":
              if training:
                  area_mean += area_std * tf.random_normal(
                      tf.shape(area_std)
                  )
              feature_concat = area_mean + size_embed
          else:
              raise ValueError("Unsupported area key mode=%s" % mode)
          feature_hidden = tf.layers.dense(
              inputs=feature_concat, units=depth, activation=tf.nn.relu
          )
          area_key = tf.layers.dense(feature_hidden, units=depth)
          return area_key
  
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 6 for calling a global function with example usage: compute_area_key(features, max_area_width, max_area_height, height, mode, training, name) and returns: area_mean, area_key, area_max, area_mean