code: def halo_exchange(x, blocks_dim, block_size_dim, halo_size, wrap=False):
    
    if halo_size == 0:
      return x
  
    block_size = block_size_dim.size
    partial_size = halo_size % block_size
    num_complete_blocks = halo_size // block_size
    parts = [x]
  
    for i in xrange(1, num_complete_blocks + 1):
      parts = ([shift(x, i, blocks_dim, wrap)] + parts +
               [shift(x, -i, blocks_dim, wrap)])
    if partial_size > 0:
      left_margin = mtf_slice(x, 0, partial_size, block_size_dim.name)
      right_margin = mtf_slice(
          x, block_size_dim.size - partial_size, partial_size,
          block_size_dim.name)
      parts = (
          [shift(right_margin, num_complete_blocks + 1, blocks_dim, wrap)]
          + parts +
          [shift(left_margin, -(num_complete_blocks + 1), blocks_dim, wrap)])
    return concat(parts, block_size_dim.name)
docs: Concat each block with the margins of adjacent blocks.
  
    Get left and right blocks_dim and concatenate along block_size_dim.
  
    Args:
      x: a Tensor.
      blocks_dim: a Dimension in x.shape
      block_size_dim: a Dimension in x.shape
      halo_size: an integer
      wrap: a boolean
  
    Returns:
      a Tensor with the same shape as x, other than in block_size_dim, whose
      size is increased by 2*halo_size.