code: def local_2d_halo_exchange(k, v, num_h_blocks, h_dim,
                             num_w_blocks, w_dim, mask_right):
    
    for blocks_dim, block_size_dim, halo_size in [
        (num_h_blocks, h_dim, h_dim.size),
        (num_w_blocks, w_dim, w_dim.size)]:
      # shape of k is [num_h_blocks, num_w_blocks, h_dim, w_dim, kv_channels]
      if halo_size > 0:
        if blocks_dim is not None:
          if mask_right:
            k = mtf.left_halo_exchange(k, blocks_dim, block_size_dim, halo_size)
            v = mtf.left_halo_exchange(v, blocks_dim, block_size_dim, halo_size)
          else:
            k = mtf.halo_exchange(k, blocks_dim, block_size_dim, halo_size)
            v = mtf.halo_exchange(v, blocks_dim, block_size_dim, halo_size)
        else:
          if mask_right:
            k = mtf.pad(k, [halo_size, None], block_size_dim.name)
            v = mtf.pad(v, [halo_size, None], block_size_dim.name)
          else:
            k = mtf.pad(k, [halo_size, halo_size], block_size_dim.name)
            v = mtf.pad(v, [halo_size, halo_size], block_size_dim.name)
    return k, v
docs: Halo exchange for keys and values for Local 2D attention.