code: def conv_with_spatial_partition(
      x,
      sampled_2d_slices,
      image_nx_dim,
      image_ny_dim,
      n_filters,
      keep_p,
      with_batch_norm,
      is_training,
      odim_name,
      variable_dtype,
      name,
  ):
      """Conv with spatial partition, batch_noram and activation."""
      if sampled_2d_slices:
          x = mtf.layers.conv2d_with_blocks(
              x,
              mtf.Dimension(odim_name, n_filters),
              filter_size=(3, 3),
              strides=(1, 1),
              padding="SAME",
              h_blocks_dim=image_nx_dim,
              w_blocks_dim=image_ny_dim,
              variable_dtype=variable_dtype,
              name=name,
          )
      else:
          x = mtf.layers.conv3d_with_blocks(
              x,
              mtf.Dimension(odim_name, n_filters),
              filter_size=(3, 3, 3),
              strides=(1, 1, 1),
              padding="SAME",
              d_blocks_dim=image_nx_dim,
              h_blocks_dim=image_ny_dim,
              variable_dtype=variable_dtype,
              name=name,
          )
  
      if with_batch_norm:
          x, bn_update_ops = mtf.layers.batch_norm(
              x,
              is_training,
              momentum=0.90,
              epsilon=0.000001,
              dims_idx_start=0,
              dims_idx_end=-1,
              name=name,
          )
      else:
          bn_update_ops = []
  
      x = mtf.leaky_relu(x, 0.1)
      x = mtf.dropout(x, is_training, keep_p)
  
      return x, bn_update_ops
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import os
  import mesh_tensorflow
  import numpy
  from six.moves import range
  import tensorflow.compat.v1
  from mesh_tensorflow.experimental import data_aug_lib
  from tensorflow.python.platform import flags
label_desc: Use image node_id 4 for calling a global function with example usage: conv_with_spatial_partition(x, sampled_2d_slices, image_nx_dim, image_ny_dim, n_filters, keep_p, with_batch_norm, is_training, odim_name, variable_dtype, name) and returns: x,  bn_update_ops