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
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]