def METHOD_NAME(have_nan_gradients):
  """Returns a function that optionally has NaN gradients.
  This serves as a hook to introduce NaN gradients to a model. This returns an
  identity function. The identity's gradient function will check if the boolean
  tensor `have_nan_gradients` is True. If so, the gradient will be NaN.
  Otherwise, the gradient will also be the identity.
  Args:
    have_nan_gradients: A scalar boolean tensor. If True, gradients will be NaN.
      Otherwise, the gradient function is the identity function.
  Returns:
    An identity function whose gradient function will return NaNs, if
    `have_nan_gradients` is True.
  """
  @custom_gradient.custom_gradient
  def identity_with_nan_gradients(x):
    """Function whose gradient is NaN iff `have_nan_gradients` is True."""
    x = array_ops.identity(x)
    def grad(dx):
      # We need this control dependency, because otherwise the NaN could be
      # produced before `dx`. This in turn could cause the final gradient to be
      # produced because `dx`, causing the loss scale to be updated before `dx`,
      # which can cause `tf.assert_equal`s to fail.
      with ops.control_dependencies([dx]):
        nan_scalar = constant_op.constant(float('NaN'), dtype=dx.dtype)
      return control_flow_ops.cond(
          have_nan_gradients,
          lambda: array_ops.fill(array_ops.shape(dx), nan_scalar),
          lambda: dx
      )
    return x, grad
  return identity_with_nan_gradients