Calculate unmasked dot-product local self-attention 2d on tpu.

  Args:
    q: a Tensor with shape [batch, heads, height, width, depth].
    k: a Tensor with shape [batch, heads, height, width, depth].
    v: a Tensor with shape [batch, heads, height, width, depth].
    bias: bias Tensor.
    max_relative_position: an integer the max relative embedding considered.
      Changing this invalidates checkpoints.
    query_shape: a two tuple indicating query shape
    dropout_rate: a floating point number.
    image_shapes: optional tuple of integer scalars.
    name: an optional string.
    make_image_summary: Whether to make an attention image summary.
    dropout_broadcast_dims:  an optional list of integers less than 4
      specifying in which dimensions to broadcast the dropout decisions.
      saves memory.

  Returns:
    [batch, heads, height, width, depth] tensor, the output of attention.

def dot_product_unmasked_attention_local_2d_tpu(
    q, k, v, bias, max_relative_position=None, query_shape=(8, 8),
    dropout_rate=0.0, image_shapes=None, name=None, make_image_summary=False,
    dropout_broadcast_dims=None):
  """Calculate unmasked dot-product local self-attention 2d on tpu.

  Args:
    q: a Tensor with shape [batch, heads, height, width, depth].
    k: a Tensor with shape [batch, heads, height, width, depth].
    v: a Tensor with shape [batch, heads, height, width, depth].
    bias: bias Tensor.
    max_relative_position: an integer the max relative embedding considered.
      Changing this invalidates checkpoints.
    query_shape: a two tuple indicating query shape
    dropout_rate: a floating point number.
    image_shapes: optional tuple of integer scalars.
    name: an optional string.
    make_image_summary: Whether to make an attention image summary.
    dropout_broadcast_dims:  an optional list of integers less than 4
      specifying in which dimensions to broadcast the dropout decisions.
      saves memory.

  Returns:
    [batch, heads, height, width, depth] tensor, the output of attention.

  """
  if max_relative_position:
    raise ValueError("Relative local 2d attention not implemented")

  with tf.variable_scope(
      name,
      default_name="dot_product_unmasked_attention_local_2d_tpu",
      values=[q, k, v]):

    # This calculation only works for self attention.
    # q, k and v must therefore have the same shape.
    q.get_shape().assert_is_compatible_with(k.get_shape())
    q.get_shape().assert_is_compatible_with(v.get_shape())
    orig_q_shape = common_layers.shape_list(q)
    # Pad query, key, value to ensure multiple of corresponding lengths.
    memory_flange = [int(query_shape[0]//2), int(query_shape[1]//2)]
    q = pad_to_multiple_2d(q, query_shape)
    k = pad_to_multiple_2d(k, query_shape)
    v = pad_to_multiple_2d(v, query_shape)
    q_shape = common_layers.shape_list(q)
    (height, width) = (q_shape[2],
                       q_shape[3])
    _, num_heads, height, width, depth_k = common_layers.shape_list(k)
    depth_v = common_layers.shape_list(v)[-1]
    num_h_blocks = height//query_shape[0]
    num_w_blocks = width//query_shape[1]
    # Extract center queries, keys, and values
    q = tf.reshape(q, [-1, height, width, depth_k])
    queries = _extract_blocks(
        q, query_shape[0], query_shape[1])
    k = tf.reshape(k, [-1, height, width, depth_k])
    keys = get_2d_local_memory_v2(
        k, query_shape, memory_flange)
    v = tf.reshape(v, [-1, height, width, depth_v])
    values = get_2d_local_memory_v2(
        v, query_shape, memory_flange)
    memory_h = query_shape[0] + 2*memory_flange[0]
    memory_w = query_shape[1] + 2*memory_flange[1]
    queries = tf.reshape(queries, [-1, num_heads, num_h_blocks, num_w_blocks,
                                   query_shape[0]*query_shape[1], depth_k])
    keys = tf.reshape(keys, [-1, num_heads, num_h_blocks, num_w_blocks,
                             memory_h*memory_w, depth_k])
    values = tf.reshape(values, [-1, num_heads, num_h_blocks, num_w_blocks,
                                 memory_h*memory_w, depth_v])
    logits = tf.matmul(queries, keys, transpose_b=True)
    if bias is not None:
      logits += bias

    weights = tf.nn.softmax(logits, name="attention_weights")
    # Dropping out the attention links for each of the heads
    weights = common_layers.dropout_with_broadcast_dims(
        weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
    if common_layers.should_generate_summaries() and make_image_summary:
      attention_image_summary(weights, image_shapes)
    ret = tf.matmul(weights, values)
    # we need to get it back to shape [batch, heads, height, width]
    ret = tf.reshape(ret, [-1, num_heads, num_h_blocks, num_w_blocks,
                           query_shape[0], query_shape[1], depth_v])
    ret = tf.transpose(ret, [0, 1, 2, 4, 3, 5, 6])
    ret = tf.reshape(ret, [-1, num_heads, num_h_blocks*query_shape[0],
                           num_w_blocks*query_shape[1], depth_v])
    # slice if padding was introduced
    ret = tf.slice(ret, [0, 0, 0, 0, 0], [-1, -1, orig_q_shape[2],
                                          orig_q_shape[3], -1])
    return ret