code: def checkpoint_test_helper(
      dvar_test,
      distribution,
      synchronization,
      aggregation,
      enable_async_ckpt,
  ):
      # This method is added since `testCheckpointing` cannot be parameterized after
      # the entire class is parameterized.
      with distribution.scope():
          v = variables_lib.Variable(
              constant_op.constant([1.0, 2.0, 3.0, 4]),
              synchronization=synchronization,
              aggregation=aggregation,
          )
  
      dvar_test.evaluate(v.initializer)
      before_save = dvar_test.evaluate(v.read_value())
  
      # Save random weights into checkpoint.
      checkpoint = trackable_utils.Checkpoint(v=v)
      ckpt_options = checkpoint_options.CheckpointOptions(
          experimental_enable_async_checkpoint=enable_async_ckpt
      )
      prefix = os.path.join(dvar_test.get_temp_dir(), "ckpt")
      with dvar_test.test_session():
          save_path = checkpoint.save(
              file_prefix=prefix, options=ckpt_options
          )
  
      # Assign inverted value.
      dvar_test.evaluate(
          v.assign(constant_op.constant([4.0, 3.0, 2.0, 1.0]))
      )
      after_assign = dvar_test.evaluate(v.read_value())
      dvar_test.assertNotAllClose(before_save, after_assign)
  
      # Restore from the checkpoint.
      with dvar_test.test_session():
          checkpoint.restore(
              save_path
          ).assert_consumed().run_restore_ops()
      after_restore = dvar_test.evaluate(v)
      dvar_test.assertAllClose(before_save, after_restore)
  
      # Another round of saving/restoring to ensure that the logic of
      # _copy_trackable_to_cpu works when a copy is already created in object_map.
      dvar_test.evaluate(
          v.assign(constant_op.constant([5.0, 6.0, 7.0, 8.0]))
      )
      before_save_1 = dvar_test.evaluate(v.read_value())
      dvar_test.assertNotAllClose(before_save_1, after_restore)
      with dvar_test.test_session():
          save_path = checkpoint.save(
              file_prefix=prefix, options=ckpt_options
          )
  
      dvar_test.evaluate(
          v.assign(constant_op.constant([8.0, 7.0, 6.0, 5.0]))
      )
      after_assign_1 = dvar_test.evaluate(v.read_value())
      dvar_test.assertNotAllClose(before_save_1, after_assign_1)
  
      with dvar_test.test_session():
          checkpoint.restore(
              save_path
          ).assert_consumed().run_restore_ops()
      after_restore_1 = dvar_test.evaluate(v)
      dvar_test.assertAllClose(before_save_1, after_restore_1)
  
imports: import copy
  import os
  from absl.testing import parameterized
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_options
  from tensorflow.python.distribute import collective_all_reduce_strategy
  from tensorflow.python.distribute import combinations
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import distribute_utils
  from tensorflow.python.distribute import packed_distributed_variable
  from tensorflow.python.distribute import parameter_server_strategy
  from tensorflow.python.distribute import ps_values
  from tensorflow.python.distribute import strategy_combinations
  from tensorflow.python.distribute import test_util
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.distribute import values
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import test
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import indexed_slices
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import check_ops
  from tensorflow.python.ops import control_flow_assert
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops import variables
  from tensorflow.python.saved_model import save
  from tensorflow.python.saved_model import save_context
  from tensorflow.python.saved_model import save_options
  from tensorflow.python.types import core
label_desc: Use image node_id 4 for calling a global function with example usage: checkpoint_test_helper(dvar_test, distribution, synchronization, aggregation, enable_async_ckpt) without return types