code: def gather_dilated_memory_blocks(x,
                                   num_memory_blocks,
                                   gap_size,
                                   query_block_size,
                                   memory_block_size,
                                   gather_indices,
                                   direction="left"):
    
    gathered_blocks = []
    # gathering memory blocks
    for block_id in range(num_memory_blocks):
      block_end_index = -(query_block_size + gap_size *
                          (block_id + 1) + memory_block_size * block_id)
      block_start_index = (
          (memory_block_size + gap_size) * (num_memory_blocks - (block_id + 1)))
      if direction != "left":
        [block_end_index,
         block_start_index] = [-block_start_index, -block_end_index]
      if block_end_index == 0:
        x_block = x[block_start_index:]
      else:
        x_block = x[block_start_index:block_end_index]
  
      def gather_dilated_1d_blocks(x, gather_indices):
        x_new = tf.gather(x, gather_indices)
        # [batch, heads, blocks, block_length, dim]
        return tf.transpose(x_new, [2, 3, 0, 1, 4])
  
      gathered_blocks.append(gather_dilated_1d_blocks(x_block, gather_indices))
    return tf.concat(gathered_blocks, 3)
docs: Gathers blocks with gaps in between.
  
    Args:
      x: Tensor of shape [length, batch, heads, depth]
      num_memory_blocks: how many memory blocks to look in "direction". Each will
        be separated by gap_size.
      gap_size: an integer indicating the gap size
      query_block_size: an integer indicating size of query block
      memory_block_size: an integer indicating the size of a memory block.
      gather_indices: The indices to gather from.
      direction: left or right
  
    Returns:
      Tensor of shape [batch, heads, blocks, block_length, depth]