code: def async_checkpoint_test_helper(test_case, x):
      # First assign an initial value 123 and save it to checkpoint.
      test_case.evaluate(x.assign(123.0))
      checkpoint = trackable_utils.Checkpoint(x=x)
      ckpt_options = checkpoint_options.CheckpointOptions(
          experimental_enable_async_checkpoint=True
      )
      prefix = os.path.join(test_case.get_temp_dir(), "ckpt")
      save_path = checkpoint.save(prefix, options=ckpt_options)
  
      # Then we modify the value to 234, restore from checkpoint, and see that the
      # value goes back to 123.
      test_case.evaluate(x.assign(234.0))
      test_case.assertNotAllClose(123.0, x.read_value())
      checkpoint.restore(save_path).assert_consumed().run_restore_ops()
      test_case.assertEqual(test_case.evaluate(x), 123.0)
  
      # 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.
      test_case.evaluate(x.assign(345.0))
      save_path = checkpoint.save(prefix, options=ckpt_options)
      test_case.evaluate(x.assign(456.0))
      test_case.assertNotAllClose(345.0, x.read_value())
      checkpoint.restore(save_path).assert_consumed().run_restore_ops()
      test_case.assertEqual(test_case.evaluate(x), 345.0)
  
file_path: tensorflow/tensorflow/python/distribute/ps_values_test.py
imports: 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 combinations
  from tensorflow.python.distribute import ps_values
  from tensorflow.python.distribute import strategy_combinations
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import test
  from tensorflow.python.ops import variable_v1
  from tensorflow.python.ops import variables
label_desc: Use image node_id 1 for calling a global function with example usage: async_checkpoint_test_helper(test_case, x) without return types