Helper function for local 2d attention.

  Args:
    x: a [batch, height, width, depth] tensor
    block_h: An integer. block height
    block_w: An inteter. block width

  returns:
    a [batch, num_heads, height/block_h, width/block_w, depth] tensor

def _extract_blocks(x, block_h, block_w):
  """Helper function for local 2d attention.

  Args:
    x: a [batch, height, width, depth] tensor
    block_h: An integer. block height
    block_w: An inteter. block width

  returns:
    a [batch, num_heads, height/block_h, width/block_w, depth] tensor
  """
  (_, height, width, depth) = common_layers.shape_list(x)
  assert height % block_h == 0
  assert width % block_w == 0
  x = tf.reshape(x, [-1, height//block_h, block_h,
                     width//block_w, block_w, depth])
  return tf.transpose(x, [0, 1, 3, 2, 4, 5])