code: def run_ncf_custom_training(
      params,
      autodist,
      keras_model,
      optimizer,
      callbacks,
      train_input_dataset,
      eval_input_dataset,
      num_train_steps,
      num_eval_steps,
      generate_input_online=True,
      return_simulation=False,
  ):
      """Runs custom training loop.
  
      Args:
          params: Dictionary containing training parameters.
          strategy: Distribution strategy to be used for distributed training.
          keras_model: Model used for training.
          optimizer: Optimizer used for training.
          callbacks: Callbacks to be invoked between batches/epochs.
          train_input_dataset: tf.data.Dataset used for training.
          eval_input_dataset: tf.data.Dataset used for evaluation.
          num_train_steps: Total number of steps to run for training.
          num_eval_steps: Total number of steps to run for evaluation.
          generate_input_online: Whether input data was generated by data producer.
              When data is generated by data producer, then train dataset must be
              re-initialized after every epoch.
  
      Returns:
          A tuple of train loss and a list of training and evaluation results.
      """
      loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
          reduction="sum", from_logits=True
      )
      train_input_iterator = tf.compat.v1.data.make_one_shot_iterator(
          train_input_dataset
      ).get_next()
  
      def step_fn(features):
          softmax_logits = keras_model(features)
          softmax_logits = tf.cast(softmax_logits, "float32")
          labels = features[rconst.TRAIN_LABEL_KEY]
          loss = loss_object(
              labels,
              softmax_logits,
              sample_weight=features[rconst.VALID_POINT_MASK],
          )
          loss *= 1.0 / params["batch_size"]
          if FLAGS.dtype == "fp16":
              loss = optimizer.get_scaled_loss(loss)
  
          grads = tf.gradients(loss, keras_model.trainable_variables)
          if FLAGS.dtype == "fp16":
              grads = optimizer.get_unscaled_gradients(grads)
          grads_and_vars = list(
              zip(grads, keras_model.trainable_variables)
          )
          if FLAGS.dense_gradient:
              grads_and_vars = neumf_model.sparse_to_dense_grads(
                  grads_and_vars
              )
          train_op = optimizer.apply_gradients(grads_and_vars)
          return (
              loss,
              train_op,
              keras_model.trainable_variables,
              grads,
              optimizer,
          )
  
      for callback in callbacks:
          callback.on_train_begin()
  
      if FLAGS.ml_perf:
          eval_summary_writer, train_summary_writer = None, None
      else:
          summary_dir = os.path.join(FLAGS.model_dir, "summaries")
          eval_summary_writer = tf.summary.create_file_writer(
              os.path.join(summary_dir, "eval")
          )
          train_summary_writer = tf.summary.create_file_writer(
              os.path.join(summary_dir, "train")
          )
  
      loss_op, train_op, vars, grads, optimizer = step_fn(
          train_input_iterator
      )
      #####################################################################
      # Create distributed session.
      #   Instead of using the original TensorFlow session for graph execution,
      #   let's use AutoDist's distributed session, in which a computational
      #   graph for distributed training is constructed.
      #
      # [original line]
      # >>> sess = tf.compat.v1.Session()
      sess = autodist.create_distributed_session()
      #####################################################################
      for epoch in range(FLAGS.num_epochs):
          for cb in callbacks:
              cb.on_epoch_begin(epoch)
  
          train_loss = 0
          for step in range(FLAGS.trial_steps):
              current_step = step + epoch * num_train_steps
              for c in callbacks:
                  c.on_batch_begin(current_step)
  
              loss_val, _ = sess.run([loss_op, train_op])
              train_loss += loss_val
  
              if train_summary_writer and step % 1000 == 0:
                  with train_summary_writer.as_default():
                      tf.summary.scalar(
                          "training_loss",
                          train_loss / (step + 1),
                          step=current_step,
                      )
  
              for c in callbacks:
                  c.on_batch_end(current_step)
  
          train_loss /= num_train_steps
          logging.info(
              "Done training epoch %s, epoch loss=%s.",
              epoch + 1,
              train_loss,
          )
  
      for c in callbacks:
          c.on_train_end()
  
      return train_loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import json
  import yaml
  import os
  import sys
  import tensorflow
  from absl import app
  from absl import flags
  from absl import logging
  from os.path import expanduser
  from tensorflow_addons.optimizers import LazyAdam
  from utils.recommendation import constants
  from utils.recommendation import movielens
  from utils.recommendation import ncf_common
  from utils.recommendation import ncf_input_pipeline
  from utils.recommendation import neumf_model
  from utils.flags import core
  from utils.logs import logger
  from utils.logs import mlperf_helper
  from utils.misc import keras_utils
  from utils.misc import model_helpers
  from autodist import AutoDist
  from autodist.strategy.ps_strategy import PS
  from autodist.strategy.ps_lb_strategy import PSLoadBalancing
  from autodist.strategy.partitioned_ps_strategy import PartitionedPS
  from autodist.strategy.all_reduce_strategy import AllReduce
  from autodist.strategy.parallax_strategy import Parallax
label_desc: Use image node_id 3 for calling a global function with example usage: run_ncf_custom_training(params, autodist, keras_model, optimizer, callbacks, train_input_dataset, eval_input_dataset, num_train_steps, num_eval_steps, generate_input_online, return_simulation) and returns: train_loss, loss,  train_op,  keras_model,  grads,  optimizer