code: def toy_model(features, mesh):
      """A toy model implemented by mesh tensorlfow."""
      batch_dim = mtf.Dimension("batch", FLAGS.batch_size)
      io_dim = mtf.Dimension("io", FLAGS.io_size)
  
      master_dtype = tf.as_dtype(FLAGS.master_dtype)
      slice_dtype = tf.as_dtype(FLAGS.slice_dtype)
      activation_dtype = tf.as_dtype(FLAGS.activation_dtype)
  
      x = mtf.import_tf_tensor(
          mesh, features, mtf.Shape([batch_dim, io_dim])
      )
      x = mtf.cast(x, activation_dtype)
      h = x
      for lnum in range(1, FLAGS.num_hidden_layers + 2):
          if lnum + 1 == FLAGS.num_hidden_layers + 2:
              # output layer
              dim = io_dim
          elif lnum % 2 == 0:
              dim = mtf.Dimension("hidden_even", FLAGS.hidden_size)
          else:
              dim = mtf.Dimension("hidden_odd", FLAGS.hidden_size)
          h = mtf.layers.dense(
              h,
              dim,
              use_bias=False,
              master_dtype=master_dtype,
              slice_dtype=slice_dtype,
              name="layer_%d" % lnum,
          )
      y = h
      g = tf.train.get_global_step()
      if FLAGS.step_with_nan >= 0:
          # Trigger NaN in the forward pass, this is used for testing whether
          # MeshTensorFlow can handle occasional NaN value.
          y += mtf.import_tf_tensor(
              mesh,
              tf.divide(
                  0.0,
                  tf.cond(
                      tf.equal(g, FLAGS.step_with_nan),
                      lambda: 0.0,
                      lambda: 1.0,
                  ),
              ),
              mtf.Shape([]),
          )
  
      loss = mtf.reduce_mean(mtf.square(y - x))
      return y, loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  import numpy
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from tensorflow.python.data.ops.dataset_ops import Dataset
  from tensorflow.python.platform import flags
  from tensorflow.python.platform import tf_logging
  from tensorflow_estimator.python.estimator import estimator
  from tensorflow_estimator.python.estimator.tpu import tpu_config
  from tensorflow_estimator.python.estimator.tpu import tpu_estimator
label_desc: Use image node_id 1 for calling a global function with example usage: toy_model(features, mesh) and returns: y,  loss