code: def toy_model(features, mesh):
    
    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 xrange(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., lambda: 1.)),
          mtf.Shape([]))
  
    loss = mtf.reduce_mean(mtf.square(y - x))
    return y, loss
docs: A toy model implemented by mesh tensorlfow.