code: def neumf_model_fn(features, labels, mode, params):
      """Model Function for NeuMF estimator."""
      if params.get("use_seed"):
          tf.set_random_seed(stat_utils.random_int32())
  
      users = features[movielens.USER_COLUMN]
      items = features[movielens.ITEM_COLUMN]
  
      user_input = tf.keras.layers.Input(tensor=users)
      item_input = tf.keras.layers.Input(tensor=items)
      logits = construct_model(user_input, item_input, params).output
  
      # Softmax with the first column of zeros is equivalent to sigmoid.
      softmax_logits = ncf_common.convert_to_softmax_logits(logits)
  
      if mode == tf.estimator.ModeKeys.EVAL:
          duplicate_mask = tf.cast(
              features[rconst.DUPLICATE_MASK], tf.float32
          )
          return _get_estimator_spec_with_metrics(
              logits,
              softmax_logits,
              duplicate_mask,
              params["num_neg"],
              params["match_mlperf"],
              use_tpu_spec=params["use_xla_for_gpu"],
          )
  
      elif mode == tf.estimator.ModeKeys.TRAIN:
          labels = tf.cast(labels, tf.int32)
          valid_pt_mask = features[rconst.VALID_POINT_MASK]
  
          mlperf_helper.ncf_print(
              key=mlperf_helper.TAGS.OPT_NAME, value="adam"
          )
          mlperf_helper.ncf_print(
              key=mlperf_helper.TAGS.OPT_LR,
              value=params["learning_rate"],
          )
          mlperf_helper.ncf_print(
              key=mlperf_helper.TAGS.OPT_HP_ADAM_BETA1,
              value=params["beta1"],
          )
          mlperf_helper.ncf_print(
              key=mlperf_helper.TAGS.OPT_HP_ADAM_BETA2,
              value=params["beta2"],
          )
          mlperf_helper.ncf_print(
              key=mlperf_helper.TAGS.OPT_HP_ADAM_EPSILON,
              value=params["epsilon"],
          )
  
          optimizer = tf.compat.v1.train.AdamOptimizer(
              learning_rate=params["learning_rate"],
              beta1=params["beta1"],
              beta2=params["beta2"],
              epsilon=params["epsilon"],
          )
          if params["use_tpu"]:
              # TODO(seemuch): remove this contrib import
              optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)
  
          mlperf_helper.ncf_print(
              key=mlperf_helper.TAGS.MODEL_HP_LOSS_FN,
              value=mlperf_helper.TAGS.BCE,
          )
  
          loss = tf.compat.v1.losses.sparse_softmax_cross_entropy(
              labels=labels,
              logits=softmax_logits,
              weights=tf.cast(valid_pt_mask, tf.float32),
          )
  
          # This tensor is used by logging hooks.
          tf.identity(loss, name="cross_entropy")
  
          global_step = tf.compat.v1.train.get_global_step()
          tvars = tf.compat.v1.trainable_variables()
          gradients = optimizer.compute_gradients(
              loss, tvars, colocate_gradients_with_ops=True
          )
          gradients = sparse_to_dense_grads(gradients)
          minimize_op = optimizer.apply_gradients(
              gradients, global_step=global_step, name="train"
          )
          update_ops = tf.compat.v1.get_collection(
              tf.compat.v1.GraphKeys.UPDATE_OPS
          )
          train_op = tf.group(minimize_op, update_ops)
  
          return tf.estimator.EstimatorSpec(
              mode=mode, loss=loss, train_op=train_op
          )
  
      else:
          raise NotImplementedError
  
file_path: autodist/examples/benchmark/utils/recommendation/neumf_model.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import sys
  from six.moves import xrange
  import tensorflow
  from utils.recommendation import constants
  from utils.recommendation import movielens
  from utils.recommendation import ncf_common
  from utils.recommendation import stat_utils
  from utils.logs import mlperf_helper
label_desc: Use image node_id 2 for calling a global function with example usage: neumf_model_fn(features, labels, mode, params) and returns: _get_estimator_spec_with_metrics, tf