import tensorflow as tf

        'a',
        py_utils.WeightParams(shape=[], init=py_utils.WeightInit.Constant(0)))
    var_a = task.theta.a
    # Make a NaN gradient.
    var_grads = py_utils.NestedMap(a=(var_a, 0. * tf.log(0.)))
    has_nan_or_inf, grad_scale, final_var_grads = task.ScaleGradients(var_grads)

    with self.session():
      tf.global_variables_initializer().run()
      with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                   'is not finite'):
        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())
        self.assertTrue(tf.is_finite(final_var_grads.a[1]).eval())


class TeacherTask(base_model.BaseTask):

  @base_layer.initializer
  def __init__(self, params):
    super(TeacherTask, self).__init__(params)
    p = self.params
    with tf.variable_scope(p.name):
      self.CreateVariable('x',
                          py_utils.WeightParams(
                              shape=[], init=py_utils.WeightInit.Constant(0)))
