Dot-product area attention.

  Args:
    q: Tensor with shape [..., length_q, depth_k].
    k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
      match with q.
    v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
      match with q.
    bias: bias Tensor (see attention_bias())
    dropout_rate: a float.
    image_shapes: optional tuple of integer scalars.
      see comments for attention_image_summary()
    name: an optional string
    attention_image_summary: the callback for making image summary of attention.
    save_weights_to: an optional dictionary to capture attention weights
      for visualization; the weights tensor will be appended there under
      a string key created from the variable scope (including name).
    dropout_broadcast_dims: an optional list of integers less than rank of q.
      Specifies in which dimensions to broadcast the dropout decisions.
    max_area_width: the max width allowed for an area.
    max_area_height: the max height allowed for an area.
    memory_height: the height of the memory.
    area_key_mode: the mode for computing area keys, which can be "mean",
      "concat", "sum", "sample_concat", and "sample_sum".
    area_value_mode: the mode for computing area values, which can be either
      "mean", or "sum".
    top_k_areas: Use the top key areas for attention.
    area_temperature: the temperature for attention softmax.
    training: indicating if it is in the training mode.
  Returns:
    Tensor with shape [..., length_q, depth_v].

def dot_product_area_attention(q,
                               k,
                               v,
                               bias,
                               dropout_rate=0.0,
                               image_shapes=None,
                               name=None,
                               attention_image_summary=None,
                               save_weights_to=None,
                               dropout_broadcast_dims=None,
                               max_area_width=1,
                               max_area_height=1,
                               memory_height=1,
                               area_key_mode="mean",
                               area_value_mode="sum",
                               top_k_areas=0,
                               area_temperature=1.0,
                               training=True):
  """Dot-product area attention.

  Args:
    q: Tensor with shape [..., length_q, depth_k].
    k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
      match with q.
    v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
      match with q.
    bias: bias Tensor (see attention_bias())
    dropout_rate: a float.
    image_shapes: optional tuple of integer scalars.
      see comments for attention_image_summary()
    name: an optional string
    attention_image_summary: the callback for making image summary of attention.
    save_weights_to: an optional dictionary to capture attention weights
      for visualization; the weights tensor will be appended there under
      a string key created from the variable scope (including name).
    dropout_broadcast_dims: an optional list of integers less than rank of q.
      Specifies in which dimensions to broadcast the dropout decisions.
    max_area_width: the max width allowed for an area.
    max_area_height: the max height allowed for an area.
    memory_height: the height of the memory.
    area_key_mode: the mode for computing area keys, which can be "mean",
      "concat", "sum", "sample_concat", and "sample_sum".
    area_value_mode: the mode for computing area values, which can be either
      "mean", or "sum".
    top_k_areas: Use the top key areas for attention.
    area_temperature: the temperature for attention softmax.
    training: indicating if it is in the training mode.
  Returns:
    Tensor with shape [..., length_q, depth_v].
  """

  tf.logging.info("dot_product_area_attention: "
                  "area_h=%d, area_w=%d, mem_h=%d, "
                  "area_key_mode=%s, area_value_mode=%s, "
                  "area_temperature=%f",
                  max_area_height, max_area_width, memory_height,
                  area_key_mode, area_value_mode,
                  area_temperature)
  with tf.variable_scope(
      name, default_name="dot_product_area_attention",
      values=[q, k, v]) as scope:
    mem_shape = common_layers.shape_list(k)
    batch_size = mem_shape[0]
    head_size = mem_shape[1]
    length = mem_shape[2]
    depth = mem_shape[3]
    k_area = compute_area_key(
        tf.reshape(k, [-1, length, depth]),
        max_area_width=max_area_width,
        max_area_height=max_area_height,
        height=memory_height,
        mode=area_key_mode,
        training=training)
    if area_value_mode == "mean":
      v_area, _, _, _, _ = compute_area_features(
          tf.reshape(v, [-1, length, depth]), max_area_width=max_area_width,
          max_area_height=max_area_height, height=memory_height)
    elif area_value_mode == "max":
      v_area, _, _ = basic_pool(tf.reshape(v, [-1, length, depth]),
                                max_area_width=max_area_width,
                                max_area_height=max_area_height,
                                height=memory_height,
                                fn=tf.reduce_max)
    elif area_value_mode == "sum":
      _, _, v_area, _, _ = compute_area_features(
          tf.reshape(v, [-1, length, depth]), max_area_width=max_area_width,
          max_area_height=max_area_height, height=memory_height)
    else:
      raise ValueError("Unsupported area value mode=%s" % area_value_mode)
    k = tf.reshape(k_area, [batch_size, head_size, -1, depth])
    v = tf.reshape(v_area, [batch_size, head_size, -1, depth])
    logits = tf.matmul(q, k, transpose_b=True)  # [..., length_q, length_kv]
    if bias is not None:
      bias = common_layers.cast_like(bias, logits)
      with tf.name_scope("compute_area_att_bias", values=[bias]):
        bias_shape = common_layers.shape_list(bias)
        mem_length = bias_shape[-1]
        bias_values = tf.reshape(
            tf.to_float(tf.less(bias, -1)), [-1, mem_length, 1])
        _, _, padding_sum, _, _ = compute_area_features(
            bias_values, max_area_width=max_area_width,
            max_area_height=max_area_height, height=memory_height)
        bias = tf.where(
            tf.cast(tf.to_int32(padding_sum), tf.bool),
            tf.fill(tf.shape(padding_sum), -np.inf),
            tf.zeros_like(padding_sum, dtype=tf.float32))
        bias = tf.reshape(bias,
                          [bias_shape[0], bias_shape[1],
                           bias_shape[2], -1])
      logits += bias
    logits = logits / area_temperature
    weights = tf.nn.softmax(logits, name="attention_weights")
    if top_k_areas > 0:
      tf.logging.info("area_attention top_k_areas=%d", top_k_areas)
      top_k = tf.minimum(common_layers.shape_list(weights)[-1], top_k_areas)
      top_weights, _ = tf.nn.top_k(weights, k=top_k)
      min_values = tf.reduce_min(top_weights, -1, keepdims=True)
      weights = tf.where(tf.greater_equal(weights, min_values),
                         weights, tf.zeros_like(weights))
      weights = tf.div(weights, tf.reduce_sum(weights, -1, keepdims=True))
    if save_weights_to is not None:
      save_weights_to[scope.name] = weights
      save_weights_to[scope.name + "/logits"] = logits
    # Drop out attention links for each head.
    weights = common_layers.dropout_with_broadcast_dims(
        weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
    if common_layers.should_generate_summaries() and attention_image_summary:
      attention_image_summary(weights, image_shapes)
    return tf.matmul(weights, v)