Strided block local self-attention.

  The 2-D sequence is divided into 2-D blocks of shape query_shape. Attention
  for a given query position can only see memory positions less than or equal to
  the query position. The memory positions are the corresponding block with
  memory_flange many positions to add to the height and width of the block
  (namely, left, top, and right).

  Args:
    q: a Tensor with shape [batch, heads, h, w, depth_k]
    k: a Tensor with shape [batch, heads, h, w, depth_k]
    v: a Tensor with shape [batch, heads, h, w, depth_v]. In the current
      implementation, depth_v must be equal to depth_k.
    query_shape: an tuple indicating the height and width of each query block.
    memory_flange: an integer indicating how much to look in height and width
      from each query block.
    name: an optional string

  Returns:
    a Tensor of shape [batch, heads, h, w, depth_v]

def local_attention_2d(q,
                       k,
                       v,
                       query_shape=(8, 16),
                       memory_flange=(8, 16),
                       name=None):
  """Strided block local self-attention.

  The 2-D sequence is divided into 2-D blocks of shape query_shape. Attention
  for a given query position can only see memory positions less than or equal to
  the query position. The memory positions are the corresponding block with
  memory_flange many positions to add to the height and width of the block
  (namely, left, top, and right).

  Args:
    q: a Tensor with shape [batch, heads, h, w, depth_k]
    k: a Tensor with shape [batch, heads, h, w, depth_k]
    v: a Tensor with shape [batch, heads, h, w, depth_v]. In the current
      implementation, depth_v must be equal to depth_k.
    query_shape: an tuple indicating the height and width of each query block.
    memory_flange: an integer indicating how much to look in height and width
      from each query block.
    name: an optional string

  Returns:
    a Tensor of shape [batch, heads, h, w, depth_v]
  """
  with tf.variable_scope(
      name, default_name="local_self_attention_2d", values=[q, k, v]):
    v_shape = common_layers.shape_list(v)

    # Pad query, key, value to ensure multiple of corresponding lengths.
    q = pad_to_multiple_2d(q, query_shape)
    k = pad_to_multiple_2d(k, query_shape)
    v = pad_to_multiple_2d(v, query_shape)
    paddings = [[0, 0], [0, 0], [memory_flange[0], memory_flange[1]],
                [memory_flange[0], memory_flange[1]], [0, 0]]
    k = tf.pad(k, paddings)
    v = tf.pad(v, paddings)

    # Set up query blocks.
    q_indices = gather_indices_2d(q, query_shape, query_shape)
    q_new = gather_blocks_2d(q, q_indices)

    # Set up key and value blocks.
    memory_shape = (query_shape[0] + 2 * memory_flange[0],
                    query_shape[1] + 2 * memory_flange[1])
    k_and_v_indices = gather_indices_2d(k, memory_shape, query_shape)
    k_new = gather_blocks_2d(k, k_and_v_indices)
    v_new = gather_blocks_2d(v, k_and_v_indices)

    attention_bias = tf.expand_dims(
        tf.to_float(embedding_to_padding(k_new)) * -1e9, axis=-2)
    output = dot_product_attention(
        q_new,
        k_new,
        v_new,
        attention_bias,
        dropout_rate=0.,
        name="local_2d",
        make_image_summary=False)
    # Put representations back into original shapes.
    padded_q_shape = common_layers.shape_list(q)
    output = scatter_blocks_2d(output, q_indices, padded_q_shape)

    # Remove the padding if introduced.
    output = tf.slice(output, [0, 0, 0, 0, 0],
                      [-1, -1, v_shape[2], v_shape[3], -1])
    return output