Creates one layer of blocks for the ResNet model.

  Args:
    inputs: `Tensor` of size `[batch, channels, height, width]`.
    filters: `int` number of filters for the first convolution of the layer.
    blocks: `int` number of blocks contained in the layer.
    strides: `int` stride to use for the first convolution of the layer. If
        greater than 1, this layer will downsample the input.
    is_training: `bool` for whether the model is training.
    name: `str`name for the Tensor output of the block layer.
    row_blocks_dim: a mtf.Dimension, row dimension which is
        spatially partitioned along mesh axis
    col_blocks_dim: a mtf.Dimension, row dimension which is
        spatially partitioned along mesh axis

  Returns:
    The output `Tensor` of the block layer.

def block_layer(inputs,
                filters,
                blocks,
                strides,
                is_training,
                name,
                row_blocks_dim=None,
                col_blocks_dim=None):
  """Creates one layer of blocks for the ResNet model.

  Args:
    inputs: `Tensor` of size `[batch, channels, height, width]`.
    filters: `int` number of filters for the first convolution of the layer.
    blocks: `int` number of blocks contained in the layer.
    strides: `int` stride to use for the first convolution of the layer. If
        greater than 1, this layer will downsample the input.
    is_training: `bool` for whether the model is training.
    name: `str`name for the Tensor output of the block layer.
    row_blocks_dim: a mtf.Dimension, row dimension which is
        spatially partitioned along mesh axis
    col_blocks_dim: a mtf.Dimension, row dimension which is
        spatially partitioned along mesh axis

  Returns:
    The output `Tensor` of the block layer.
  """
  with tf.variable_scope(name, default_name="block_layer"):
    # Only the first block per block_layer uses projection_shortcut and strides
    def projection_shortcut(inputs, kernel):
      """Project identity branch."""
      inputs = mtf.conv2d_with_blocks(
          inputs,
          kernel,
          strides=strides,
          padding="SAME",
          h_blocks_dim=None, w_blocks_dim=col_blocks_dim)
      return batch_norm_relu(
          inputs, is_training, relu=False)

    inputs = bottleneck_block(
        inputs,
        filters,
        is_training,
        strides=strides,
        projection_shortcut=projection_shortcut,
        row_blocks_dim=row_blocks_dim,
        col_blocks_dim=col_blocks_dim)

    for i in range(1, blocks):
      with tf.variable_scope("bottleneck_%d" % i):
        inputs = bottleneck_block(
            inputs,
            filters,
            is_training,
            strides=[1, 1, 1, 1],
            projection_shortcut=None,
            row_blocks_dim=row_blocks_dim,
            col_blocks_dim=col_blocks_dim)

    return inputs