code: def main(autodist):
      TRUE_W = 3.0
      TRUE_b = 2.0
      NUM_EXAMPLES = 1000
      EPOCHS = 1
  
      # For Integration Value Test: (For more information, check the assertions below)
      seed = 456 if bool(ENV.AUTODIST_WORKER.val) else 123
      np.random.seed(seed)
  
      inputs = np.random.randn(NUM_EXAMPLES)
      noises = np.random.randn(NUM_EXAMPLES)
      outputs = inputs * TRUE_W + TRUE_b + noises
  
      class MyIterator:
          def initialize(self):
              return tf.zeros(1)
  
          def get_next(self):
              # a fake one
              return inputs
  
      inputs_iterator = MyIterator()
      with tf.Graph().as_default(), autodist.scope():
          x = tf.compat.v1.placeholder(shape=[None], dtype=tf.float32)
          y = tf.compat.v1.placeholder(shape=[None], dtype=tf.float32)
  
          W = tf.Variable(5.0, name="W")
          b = tf.Variable(0.0, name="b")
  
          def train_step(x):
              def f(x):
                  return W * x + b
  
              def l(predicted_y, desired_y):
                  return tf.reduce_mean(
                      tf.square(predicted_y - desired_y)
                  )
  
              major_version, _, _ = tf.version.VERSION.split(".")
              if major_version == "1":
                  optimizer = tf.train.GradientDescentOptimizer(0.01)
              else:
                  optimizer = tf.optimizers.SGD(0.01)
  
              with tf.GradientTape() as tape:
                  loss = l(f(x), y)
                  vs = [W, b]
  
                  # gradients = tape.gradient(target=loss, sources=vs)
                  gradients = tf.gradients(loss, vs)
  
                  train_op = optimizer.apply_gradients(
                      zip(gradients, vs)
                  )
              return loss, train_op, b
  
          assert EPOCHS == 1
          fetches = train_step(x)
          saver = Saver([W, b])
          session = autodist.create_distributed_session()
          for epoch in range(EPOCHS):
              l_val, _, _ = session.run(
                  fetches=fetches,
                  options=tf.compat.v1.RunOptions(
                      trace_level=tf.compat.v1.RunOptions.FULL_TRACE
                  ),
                  feed_dict={x: inputs_iterator.get_next(), y: outputs},
              )
              print("loss:", l_val)
              # Seperate the fetches of var to guarantee the state
              W_val, b_val = session.run([W, b])
  
          # Try to save the two variables
          checkpoint_dir = "/tmp/ckpt/"
          if not os.path.exists(checkpoint_dir):
              os.mkdir(checkpoint_dir)
          checkpoint_suffix = "c0"
          checkpoint_name = checkpoint_dir + checkpoint_suffix
          saver.save(session, checkpoint_name, global_step=epoch)
          print("Checkpoint saved at {%s}" % checkpoint_name)
  
          num_workers = len(autodist._cluster.cluster_spec["worker"])
          if getattr(autodist._strategy_builder, "_sync", True):
              # Integration Value Test:
              # It requires np.random.seed to be 123 on chief, 456 on worker.
              # Variable b's gradients[1] == -4.17503 on chief, == -4.05530 on worker
              # When SGD learning rate == 0.01 and b is initialied by zero,
              # the updated variable b value can be verified as below
              if num_workers == 1:
                  assert np.allclose(b_val, 0.01 * 4.17503)
              elif num_workers == 2:
                  # Between-graph dense conditional accumulator average verification
  
                  def get_gpu_distribution(resource_spec):
                      """A list contains the distribution of devices."""
                      gpu_dist = dict()
                      for gpu_str, _ in resource_spec.gpu_devices:
                          ip = gpu_str.split(":")[0]
                          if ip in gpu_dist:
                              gpu_dist[ip] += 1
                          else:
                              gpu_dist[ip] = 1
                      return list(gpu_dist.values())
  
                  dist = get_gpu_distribution(autodist._resource_spec)
                  # if uneven gpu distribution, allreduce grad is a weighted sum.
                  if dist[0] != dist[1] and (
                      isinstance(autodist._strategy_builder, AllReduce)
                      or isinstance(
                          autodist._strategy_builder, Parallax
                      )
                      or isinstance(
                          autodist._strategy_builder, PartitionedAR
                      )
                      or isinstance(
                          autodist._strategy_builder,
                          RandomAxisPartitionAR,
                      )
                  ):
                      assert np.allclose(
                          b_val,
                          0.01
                          * (4.17503 * dist[0] + 4.05530 * dist[1])
                          / (dist[0] + dist[1]),
                      )
                  else:
                      assert np.allclose(
                          b_val, 0.01 * (4.17503 + 4.05530) / 2
                      )
                  # TODO: between graph sparse verification
  
          # check the checkpoint existence
          checkpoint = checkpoint_name + "-" + str(epoch)
          assert os.path.exists(checkpoint + ".meta")  # meta file
          assert os.path.exists(checkpoint + ".index")  # meta file
          assert os.path.exists(
              checkpoint + ".data-00000-of-00001"
          )  # meta file
          print("Checkpoint {} exists.".format(checkpoint))
  
          # Now check restoration and the correctness of the checkpoint
          tf_saver = tf.compat.v1.train.Saver([W, b])
          tf_sess = tf.compat.v1.Session()
          tf_saver.restore(
              tf_sess, tf.train.latest_checkpoint(checkpoint_dir)
          )
          restore_W_val, restore_b_val = tf_sess.run([W, b])
          print(
              "original value {} vs. restored value {}".format(
                  W_val, restore_W_val
              )
          )
          print(
              "original value {} vs. restored value {}".format(
                  b_val, restore_b_val
              )
          )
          assert np.allclose(b_val, restore_b_val)
          assert np.allclose(W_val, restore_W_val)
  
imports: import os
  import numpy
  import tensorflow
  from autodist.const import ENV
  from autodist.checkpoint.saver import Saver
  from autodist.strategy import AllReduce, Parallax, PartitionedAR, RandomAxisPartitionAR
label_desc: Use image node_id 1 for calling a global function with example usage: main(autodist) and returns: tf, inputs, loss,  train_op,  b, unknown, tf, list