code: def get_input_mtf_shapes(dataset_str):
      """Returns a list of mtf.Shapes of input tensors."""
      if dataset_str == "train":
          batch_dim = mtf.Dimension("batch", FLAGS.batch_size_train)
      else:
          assert dataset_str == "eval"
          batch_dim = mtf.Dimension("batch", FLAGS.batch_size_eval)
      image_nx_dim = mtf.Dimension(
          "image_nx_block", FLAGS.image_nx_block
      )
      image_ny_dim = mtf.Dimension(
          "image_ny_block", FLAGS.image_ny_block
      )
      image_sx_dim = mtf.Dimension(
          "image_sx_block", FLAGS.ct_resolution // FLAGS.image_nx_block
      )
      image_sy_dim = mtf.Dimension(
          "image_sy_block", FLAGS.ct_resolution // FLAGS.image_ny_block
      )
  
      batch_spatial_dims = [
          batch_dim,
          image_nx_dim,
          image_sx_dim,
          image_ny_dim,
          image_sy_dim,
      ]
      if not FLAGS.sampled_2d_slices:
          image_sz_dim = mtf.Dimension(
              "image_sz_block", FLAGS.ct_resolution
          )
          batch_spatial_dims += [image_sz_dim]
  
      image_c_dim = mtf.Dimension("image_c", FLAGS.image_c)
      mtf_image_shape = mtf.Shape(batch_spatial_dims + [image_c_dim])
  
      label_c_dim = mtf.Dimension("label_c", FLAGS.label_c)
      mtf_label_shape = mtf.Shape(batch_spatial_dims + [label_c_dim])
  
      return [mtf_image_shape, mtf_label_shape]
  
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 3 for calling a global function with example usage: get_input_mtf_shapes(dataset_str) and returns: list