code: def testCheckpoint(self, strategy_fn):
      strategy = strategy_fn()
      if (
          isinstance(strategy, mirrored_strategy.MirroredStrategy)
          and not context.executing_eagerly()
      ):
          # TODO(b/121381184): Enable running the test in this case.
          return
  
      with self.test_session(), strategy.scope():
          # Build and run a simple model.
          var = variables.Variable([2.0])
          loss_scale = loss_scale_module.DynamicLossScale(
              initial_loss_scale=1.0,
              increment_period=2.0,
              multiplier=2.0,
          )
          opt = momentum.MomentumOptimizer(1.0, momentum=1.0)
          opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
              opt, loss_scale
          )
          run_fn = lambda: opt.minimize(
              lambda: var + 1.0, var_list=[var]
          )
          opt_op = strategy.experimental_run(run_fn)
          self.evaluate(variables.global_variables_initializer())
          self.evaluate(opt_op)
          self.assertEqual(self.evaluate(loss_scale()), 1.0)
          self.assertEqual(self.evaluate(loss_scale._num_good_steps), 1)
  
          # Save a checkpoint.
          checkpoint = trackable_utils.Checkpoint(optimizer=opt)
          prefix = os.path.join(self.get_temp_dir(), "ckpt")
          save_path = checkpoint.save(prefix)
  
          # Run model again.
          self.evaluate(strategy.experimental_run(run_fn))
          self.assertEqual(self.evaluate(loss_scale()), 2.0)
          self.assertEqual(self.evaluate(loss_scale._num_good_steps), 0)
  
          # Load checkpoint and ensure loss scale is back to it's original value.
          status = checkpoint.restore(save_path)
          status.assert_consumed()
          status.run_restore_ops()
          self.assertEqual(self.evaluate(loss_scale()), 1.0)
          self.assertEqual(self.evaluate(loss_scale._num_good_steps), 1)
  
imports: import os
  from absl.testing import parameterized
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import mirrored_strategy
  from tensorflow.python.eager import context
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_conversion
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import check_ops
  from tensorflow.python.ops import custom_gradient
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
  from tensorflow.python.training import gradient_descent
  from tensorflow.python.training import momentum
  from tensorflow.python.training.experimental import loss_scale
  from tensorflow.python.training.experimental import loss_scale_optimizer
label_desc: Use image node_id 8 for calling the MixedPrecisionLossScaleOptimizerTest obj's underlying member method code with example usage: obj.testCheckpoint(strategy_fn) without return types