Helper function. Assumes that memory_flange is half of query sizes.

  This function splits the tensor of width 'n' into two halves, where the
  first half gets the width indices 0, 2, 4.. and the second half gets the
  width indices 3, 5, ... We also fuse two blocks along the h dimension.

  Args:
    x: a 6-d tensor.

  Returns:
    x_left_blocks, x_right_blocks: Two 6-d tensors

def _get_left_right_blocks(x):
  """Helper function. Assumes that memory_flange is half of query sizes.

  This function splits the tensor of width 'n' into two halves, where the
  first half gets the width indices 0, 2, 4.. and the second half gets the
  width indices 3, 5, ... We also fuse two blocks along the h dimension.

  Args:
    x: a 6-d tensor.

  Returns:
    x_left_blocks, x_right_blocks: Two 6-d tensors
  """
  (_, x_num_outer_h_blocks, x_num_outer_w_blocks, x_memory_flange_h,
   x_memory_flange_w, depth) = common_layers.shape_list(x)
  x_left_right_blocks = tf.slice(x,
                                 [0, 1, 0, 0, 0, 0],
                                 [-1, x_num_outer_h_blocks-2, -1, -1,
                                  -1, -1])
  num_blocks_h = (x_num_outer_h_blocks-2)//2
  x_left_right_blocks = tf.reshape(x_left_right_blocks,
                                   [-1,
                                    num_blocks_h,
                                    2, x_num_outer_w_blocks,
                                    x_memory_flange_h,
                                    x_memory_flange_w, depth])
  x_left_right_blocks = tf.transpose(x_left_right_blocks,
                                     [0, 1, 3, 2, 4, 5, 6])
  x_left_right_blocks = tf.reshape(x_left_right_blocks,
                                   [-1, num_blocks_h,
                                    x_num_outer_w_blocks, 2*x_memory_flange_h,
                                    x_memory_flange_w, depth])
  # get it ready for splitting the left and right memory blocks
  x_left_blocks, x_right_blocks = _split_along_width(x_left_right_blocks)

  return x_left_blocks, x_right_blocks