code: def testDynamicLossScale(self, strategy_fn):
      strategy = strategy_fn()
      learning_rate = 2.0
      expected_gradient = resource_variable_ops.ResourceVariable(
          learning_rate / strategy.num_replicas_in_sync
      )
      with strategy.scope():
          var = variables.Variable([5.0])
          opt = gradient_descent.GradientDescentOptimizer(learning_rate)
          loss_scale = loss_scale_module.DynamicLossScale(
              initial_loss_scale=2, increment_period=1, multiplier=2
          )
          opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
              opt, loss_scale
          )
          self.assertEqual(
              loss_scale.initial_loss_scale
              % strategy.num_replicas_in_sync,
              0,
          )
  
          run_fn = self._run_fn_with_grad_check(
              strategy, var, opt, expected_gradient
          )
          run_op = strategy.experimental_run(run_fn)
          self.evaluate(variables.global_variables_initializer())
          self._run_if_in_graph_mode(run_op)
          # The loss is the identity of the variable. Therefore the gradient is 1,
          # and so the variable will be init_val - grad * lr == 5 - 1 * 2 == 3
          self.assertAllClose([3.0], self.evaluate(var))
  
          # Loss scale will be double, so the expected gradient is also doubled.
          self.evaluate(
              expected_gradient.assign(
                  2 * learning_rate / strategy.num_replicas_in_sync
              )
          )
          run_op = strategy.experimental_run(run_fn)
          self._run_if_in_graph_mode(run_op)
          # As before, the 2 is subtracted from the variable, making it's new value
          # 1.
          self.assertAllClose([1.0], self.evaluate(var))
  
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 5 for calling the MixedPrecisionLossScaleOptimizerTest obj's underlying member method code with example usage: obj.testDynamicLossScale(strategy_fn) without return types