conv2d operation with spatial partitioning.

  Spatial partitioning is implemented by decomposing the image into blocks.
  Block dimensions represented as h_blocks_dim and w_blocks_dim can be split
  along the mesh axis. If split, then we do a halo exchange where each block
  receives the part of the image from its left and right neighbors necessary to
  do the convolution. Exchange can involve complete or partial blocks depending
  on the filter height and width.

  Currently, only "SAME" padding with dilation rate of 1 is supported.

  Args:
    conv_input: a Tensor of shape
      [batch, h_blocks_dim, w_blocks_dim, h_dim, w_dim, in_channels_dim]
    conv_filter: a Tensor of shape
      [filter_height, filter_width, in_channels_dim, out_channels_dim]
    strides: A list of ints. 1-D tensor of length 4.
    padding: string, "SAME". The type of padding algorithm to use.
      Valid is not currently supported.
    h_blocks_dim: Dimension representing number of height blocks.
    w_blocks_dim: Dimension representing number of height blocks.
    name: A name for the operation (optional).

  Returns:
    A Tensor of shape
      [batch, h_blocks_dim, w_blocks_dim, h_dim, w_dim, out_channels_dim]

def conv2d_with_blocks(
    conv_input,
    conv_filter,
    strides,
    padding,
    h_blocks_dim=None,
    w_blocks_dim=None,
    name=None):
  """conv2d operation with spatial partitioning.

  Spatial partitioning is implemented by decomposing the image into blocks.
  Block dimensions represented as h_blocks_dim and w_blocks_dim can be split
  along the mesh axis. If split, then we do a halo exchange where each block
  receives the part of the image from its left and right neighbors necessary to
  do the convolution. Exchange can involve complete or partial blocks depending
  on the filter height and width.

  Currently, only "SAME" padding with dilation rate of 1 is supported.

  Args:
    conv_input: a Tensor of shape
      [batch, h_blocks_dim, w_blocks_dim, h_dim, w_dim, in_channels_dim]
    conv_filter: a Tensor of shape
      [filter_height, filter_width, in_channels_dim, out_channels_dim]
    strides: A list of ints. 1-D tensor of length 4.
    padding: string, "SAME". The type of padding algorithm to use.
      Valid is not currently supported.
    h_blocks_dim: Dimension representing number of height blocks.
    w_blocks_dim: Dimension representing number of height blocks.
    name: A name for the operation (optional).

  Returns:
    A Tensor of shape
      [batch, h_blocks_dim, w_blocks_dim, h_dim, w_dim, out_channels_dim]
  """
  filter_h_dim, filter_w_dim = conv_filter.shape.dims[:2]
  assert filter_h_dim.size % 2 == 1
  assert filter_w_dim.size % 2 == 1
  h_dim, w_dim = conv_input.shape.dims[-3:-1]

  # If h_blocks_dim and w_blocks_dim is not split, directly call conv2d.
  if h_blocks_dim is None and w_blocks_dim is None:
    return conv2d(conv_input, conv_filter, strides, padding, name)

  # Padding 'VALID' is not supported yet.
  if padding != "SAME":
    raise NotImplementedError("conv2d_with_blocks requires padding=SAME")

  # Halo exchange for h_blocks and w_blocks.
  for blocks_dim, block_size_dim, halo_size in [
      (h_blocks_dim, h_dim, filter_h_dim.size // 2),
      (w_blocks_dim, w_dim, filter_w_dim.size // 2)]:
    if halo_size > 0:
      if blocks_dim is not None:
        conv_input = halo_exchange(
            conv_input, blocks_dim, block_size_dim, halo_size)
      else:
        conv_input = pad(
            conv_input, [halo_size, halo_size], block_size_dim.name)
  return conv2d(conv_input, conv_filter, strides, "VALID", name)