class_docstr: Test for cases in which enable_check_numerics() catches infs or nans.
class_docstr_tok: ["Test","for","cases","in","which","enable_check_numerics","(",")","catches","infs","or","nans","."]
docstr: Test that callback catches NaN in a gradient function during backprop.
docstr_tok: ["Test","that","callback","catches","NaN","in","a","gradient","function","during","backprop","."]
code: def testCustomGradientWithNaNWithTfFunction(self):
      """Test that callback catches NaN in a gradient function during backprop."""
      check_numerics_callback.enable_check_numerics()
  
      @custom_gradient.custom_gradient
      def func_with_bad_grad(x):
          output = math_ops.sin(x)
  
          @def_function.function
          def grad(dy):
              # `dy` will come in as 1.0. Taking log of -1.0 leads to NaN.
              return math_ops.log(-dy)
  
          return output, grad
  
      x = constant_op.constant(-2.0, dtype=dtypes.float16)
  
      def f(x):
          return func_with_bad_grad(x)
  
      message = self._assertRaisesInvalidArgumentErrorAndGetMessage(
          lambda: gradient_checker_v2.compute_gradient(f, [x])
      )
  
      # Check the content of the error message.
      self.assertTrue(re.search(r"graph op.*\"Log\"", message))
      self.assertTrue(re.search(r"dtype.*float16", message))
      if context.executing_eagerly():
          self.assertIn("shape: ()\n", message)
      self.assertTrue(
          re.search(r"Input tensor.*Tensor.*Neg:0", message)
      )
      self.assertIn("grad", message)
  
imports: import re
  import numpy
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.debug.lib import check_numerics_callback
  from tensorflow.python.eager import backprop
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_grad
  from tensorflow.python.ops import custom_gradient
  from tensorflow.python.ops import gen_nn_ops
  from tensorflow.python.ops import gradient_checker_v2
  from tensorflow.python.ops import math_grad
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import googletest
  from tensorflow.python.platform import test
label_desc: Use image node_id 12 for calling the CheckNumericsCallbackUnhealthyTest obj's underlying member method code with example usage: obj.testCustomGradientWithNaNWithTfFunction() and returns: output,  grad, func_with_bad_grad, math_ops