import tensorflow as tf

    FLAGS.enable_check_numerics = False
    p = self.TestParams()
    p.input = base_input_generator.BaseSequenceInputGenerator.Params()
    task = p.cls(p)
    task.CreateVariable(
        'a',
        py_utils.WeightParams(shape=[], init=py_utils.WeightInit.Constant(0)))
    var_a = task.theta.a
    # Infinite gradient.
    var_grads = py_utils.NestedMap(a=(var_a, tf.log(0.)))
    has_nan_or_inf, grad_scale, final_var_grads = task.ScaleGradients(var_grads)

    with self.session():
      tf.global_variables_initializer().run()
      self.assertTrue(has_nan_or_inf.eval())
      self.assertEqual(0., grad_scale.eval())
      # The final gradient must be finite.
      self.assertFalse(tf.is_nan(final_var_grads.a[1]).eval())
