code: def bottleneck_block(
      inputs,
      filters,
      is_training,
      strides,
      projection_shortcut=None,
      row_blocks_dim=None,
      col_blocks_dim=None,
  ):
      """Bottleneck block variant for residual networks with BN after convolutions.
  
      Args:
        inputs: a `mtf.Tensor` of shape
            `[batch_dim, row_blocks, col_blocks, rows, cols, in_channels]`.
        filters: `int` number of filters for the first two convolutions. Note
            that the third and final convolution will use 4 times as many filters.
        is_training: `bool` for whether the model is in training mode.
        strides: `int` block stride. If greater than 1, this block will ultimately
            downsample the input.
        projection_shortcut: `function` to use for projection shortcuts (typically
            a 1x1 convolution to match the filter dimensions). If None, no
            projection is used and the input is passed as unchanged through the
            shortcut connection.
        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.
      """
      shortcut = inputs
  
      if projection_shortcut is not None:
          filters_dim = mtf.Dimension("filtersp", filters)
          shortcut = projection_shortcut(inputs, filters_dim)
  
      # First conv block
      inputs = mtf.layers.conv2d_with_blocks(
          inputs,
          mtf.Dimension("filters1", filters),
          filter_size=[1, 1],
          strides=[1, 1],
          padding="SAME",
          h_blocks_dim=None,
          w_blocks_dim=col_blocks_dim,
          name="conv0",
      )
  
      # TODO(nikip): Add Dropout?
      inputs = batch_norm_relu(inputs, is_training)
  
      # Second conv block
      inputs = mtf.layers.conv2d_with_blocks(
          inputs,
          mtf.Dimension("filters2", 4 * filters),
          filter_size=[3, 3],
          strides=[1, 1],
          padding="SAME",
          h_blocks_dim=row_blocks_dim,
          w_blocks_dim=col_blocks_dim,
          name="conv1",
      )
  
      inputs = batch_norm_relu(inputs, is_training)
  
      # Third wide conv filter block
      inputs = mtf.layers.conv2d_with_blocks(
          inputs,
          mtf.Dimension("filters3", filters),
          filter_size=[1, 1],
          strides=strides,
          padding="SAME",
          h_blocks_dim=None,
          w_blocks_dim=col_blocks_dim,
          name="conv2",
      )
  
      # TODO(nikip): Althought the original resnet code has this batch norm, in our
      # setup this is causing no gradients to be passed. Investigate further.
      # inputs = batch_norm_relu(inputs, is_training, relu=True)
  
      # TODO(nikip): Maybe add residual with a projection?
      return mtf.relu(
          shortcut
          + mtf.rename_dimension(
              inputs,
              inputs.shape.dims[-1].name,
              shortcut.shape.dims[-1].name,
          )
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import mesh_tensorflow
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.utils import mtf_model
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 2 for calling a global function with example usage: bottleneck_block(inputs, filters, is_training, strides, projection_shortcut, row_blocks_dim, col_blocks_dim) and returns: mtf