Stitches together the local 2d memory blocks.

  Args:
    x: a [batch, height, width, depth tensor]
    query_shape: 2-d integer list of query shape
    memory_flange: 2-d integer list of memory flanges

  Returns:
    x: A [batch, num_h_blocks, num_w_blocks,
          query_shape[0]+2*memory_flange[0],query_shape[1]+2*memory_flange[1]]
          tensor.

def get_2d_local_memory(x, query_shape, memory_flange):
  """Stitches together the local 2d memory blocks.

  Args:
    x: a [batch, height, width, depth tensor]
    query_shape: 2-d integer list of query shape
    memory_flange: 2-d integer list of memory flanges

  Returns:
    x: A [batch, num_h_blocks, num_w_blocks,
          query_shape[0]+2*memory_flange[0],query_shape[1]+2*memory_flange[1]]
          tensor.
  """
  (_, height, width, depth_x) = common_layers.shape_list(x)
  x_center_blocks = _extract_blocks(x, query_shape[0], query_shape[1])
  # add extra padding to x so that we can extract the memory region
  # around the center
  paddings = [[0, 0], [memory_flange[0], memory_flange[0]],
              [memory_flange[1], memory_flange[1]], [0, 0]]
  padded_x = tf.pad(x, paddings)
  padded_x.set_shape([None, height+2*memory_flange[0],
                      width+2*memory_flange[1], depth_x])
  x_outer_memory_blocks = _extract_blocks(padded_x,
                                          memory_flange[0], memory_flange[1])
  # We'll extract left and right memory blocks, top and bottom memory blocks,
  # and then the corner memory blocks

  # Each of these after  will have shape
  # [batch, num_h_blocks, num_w_blocks, query_shape[0],
  # memory_flange[1], depth]
  x_left_blocks, x_right_blocks = _get_left_right_blocks(
      x_outer_memory_blocks)
  t_hw_block = lambda x: tf.transpose(x, [0, 2, 1, 4, 3, 5])
  # now to get top and bottom blocks, we should just transpose the outer
  # blocks, call the same function and transpose back to get shape
  # [batch, num_h_blocks, num_w_blocks, memory_flange[0],
  # query_shape[1], depth]
  x_top_center_blocks, x_bottom_center_blocks = (
      map(t_hw_block, _get_left_right_blocks(
          t_hw_block(x_outer_memory_blocks))))

  # now to get the corner blocks
  x_left_corner_blocks, x_right_corner_blocks = _split_along_width(
      x_outer_memory_blocks)
  # now to extract top and bottom for both k and v
  # we need to transpose because _split_along_width separates along
  # the width
  # each of these should have shape [batch, num_h_blocks,
  # num_w_blocks, memory_flange[0], memory_flange[1], depth]

  t_hw = lambda x: tf.transpose(x, [0, 2, 1, 3, 4, 5])
  x_top_left_corner_blocks, x_bottom_left_corner_blocks = (
      map(t_hw, _split_along_width(t_hw(x_left_corner_blocks))))
  x_top_right_corner_blocks, x_bottom_right_corner_blocks = (
      map(t_hw, _split_along_width(t_hw(x_right_corner_blocks))))

  # The memory is top_left     top_center    top_right
  #               left_center  middle        right_center
  #               bottom_left  bottom_center bottom_right
  # Assembling the above row by row
  # first [x_top_left, x_top, x_top_right]
  # to get [batch, num_h_blocks, num_w_blocks, memory_flange[0],
  # query_shape[1]+2*memory_flange[1], depth]
  # then [x_left, x_center, x_right]
  # then [x_bottom_left, x_bottom, x_bottom_right]
  x_top_memory = tf.concat(
      [x_top_left_corner_blocks,
       x_top_center_blocks,
       x_top_right_corner_blocks], axis=4)
  x_middle_memory = tf.concat(
      [x_left_blocks, x_center_blocks, x_right_blocks], axis=4)
  x_bottom_memory = tf.concat(
      [x_bottom_left_corner_blocks,
       x_bottom_center_blocks,
       x_bottom_right_corner_blocks], axis=4)

  # concat along height
  x = tf.concat([x_top_memory, x_middle_memory, x_bottom_memory], axis=3)
  return x