code: def create_model(
      bert_config,
      is_training,
      input_ids,
      input_mask,
      segment_ids,
      labels,
      num_labels_dim,
      layout,
      mesh_shape,
  ):
      """Creates a classification model."""
      model = bert_lib.BertModel(
          config=bert_config,
          is_training=is_training,
          input_ids=input_ids,
          input_mask=input_mask,
          token_type_ids=segment_ids,
          layout=layout,
          mesh_shape=mesh_shape,
      )
  
      # In the demo, we are doing a simple classification task on the entire
      # segment.
      #
      # If you want to use the token-level output, use model.get_sequence_output()
      # instead.
      output_layer = model.get_pooled_output()
  
      hidden_dim = output_layer.shape[-1]
      mesh = input_ids.mesh
  
      output_weights = mtf.get_variable(
          mesh,
          "output_weights",
          shape=[num_labels_dim, hidden_dim],
          initializer=tf.truncated_normal_initializer(stddev=0.02),
      )
  
      output_bias = mtf.get_variable(
          mesh,
          "output_bias",
          shape=[num_labels_dim],
          initializer=tf.zeros_initializer(),
      )
  
      with tf.variable_scope("loss"):
          # I.e., 0.1 dropout
          output_layer = mtf.dropout(
              output_layer, is_training, keep_prob=0.9
          )
          logits = mtf.einsum(
              [output_layer, output_weights], reduced_dims=[hidden_dim]
          )
          logits = logits + output_bias
          probabilities = mtf.softmax(
              logits, reduced_dim=num_labels_dim
          )
  
          per_example_loss = (
              mtf.layers.softmax_cross_entropy_with_logits(
                  logits, labels, vocab_dim=num_labels_dim
              )
          )
          loss = (
              mtf.reduce_mean(per_example_loss) + model.get_extra_loss()
          )
  
          return (loss, per_example_loss, logits, probabilities)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import csv
  import os
  import mesh_tensorflow
  import mesh_tensorflow.bert.bert
  import mesh_tensorflow.bert.optimization
  import mesh_tensorflow.bert.tokenization
  from six.moves import range
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 5 for calling a global function with example usage: create_model(bert_config, is_training, input_ids, input_mask, segment_ids, labels, num_labels_dim, layout, mesh_shape) and returns: loss,  per_example_loss,  logits,  probabilities