code: def main(
      train_dataset,
      test_dataset=None,
      tf_config=None,
      server=None,
      config=None,
  ):
      iterator = tf.data.Iterator.from_structure(
          train_dataset.output_types, train_dataset.output_shapes
      )
      next_element = iterator.get_next()
  
      train_init_op = iterator.make_initializer(train_dataset)
      test_init_op = iterator.make_initializer(test_dataset)
  
      # create feature column
      wide_column, deep_column = build_feature_columns()
  
      # create variable partitioner for distributed training
      num_ps_replicas = len(tf_config["ps_hosts"]) if tf_config else 0
      input_layer_partitioner = (
          partitioned_variables.min_max_variable_partitioner(
              max_partitions=num_ps_replicas,
              min_slice_size=config["input_layer_partitioner"] << 20,
          )
          if config["input_layer_partitioner"]
          else None
      )
      dense_layer_partitioner = (
          partitioned_variables.min_max_variable_partitioner(
              max_partitions=num_ps_replicas,
              min_slice_size=config["dense_layer_partitioner"] << 10,
          )
          if config["dense_layer_partitioner"]
          else None
      )
  
      # Session config
      sess_config = tf.ConfigProto()
      sess_config.inter_op_parallelism_threads = config["inter"]
      sess_config.intra_op_parallelism_threads = config["intra"]
  
      # Session hooks
      hooks = []
  
      if config["smartstaged"] and not config["tf"]:
          """Smart staged Feature"""
          next_element = tf.staged(
              next_element, num_threads=4, capacity=40
          )
          sess_config.graph_options.optimizer_options.do_smart_stage = (
              True
          )
          hooks.append(tf.make_prefetch_hook())
      if config["op_fusion"] and not config["tf"]:
          """Auto Graph Fusion"""
          sess_config.graph_options.optimizer_options.do_op_fusion = (
              True
          )
      if config["micro_batch"] and not config["tf"]:
          """Auto Mirco Batch"""
          sess_config.graph_options.optimizer_options.micro_batch_num = config[
              "micro_batch"
          ]
  
      # create model
      model = WDL(
          wide_column=wide_column,
          deep_column=deep_column,
          linear_learning_rate=config["linear_learning_rate"],
          deep_learning_rate=config["deep_learning_rate"],
          optimizer_type=config["optimizer"],
          bf16=config["bf16"],
          stock_tf=config["tf"],
          adaptive_emb=config["adaptive_emb"],
          inputs=next_element,
          input_layer_partitioner=input_layer_partitioner,
          dense_layer_partitioner=dense_layer_partitioner,
      )
  
      # Run model training and evaluation
      train(
          sess_config,
          hooks,
          model,
          train_init_op,
          config,
          tf_config,
          server,
      )
      # TODO: the original script won't evaluate in distributed mode, has issue in evaluation results
      if not config["no_eval"]:
          eval_acc, eval_auc = eval(
              sess_config,
              hooks,
              model,
              test_init_op,
              config["test_steps"],
              config["checkpoint_dir"],
          )
          return eval_acc, eval_auc
      return None
  
imports: import time
  import argparse
  import tensorflow
  import os
  import sys
  import math
  import collections
  from tensorflow.python.client import timeline
  import json
  from tensorflow.python.ops import partitioned_variables
  import ray
  from bigdl.dllib.utils.log4Error import *
  from bigdl.orca import init_orca_context, stop_orca_context, OrcaContext
  from bigdl.orca.data.utils import process_spark_xshards
  from bigdl.orca.learn.utils import maybe_dataframe_to_xshards
  from bigdl.friesian.feature import FeatureTable
label_desc: Use image node_id 4 for calling a global function with example usage: main(train_dataset, test_dataset, tf_config, server, config) and returns: None, eval_acc,  eval_auc