ForwardOperation forward_op;
  Status s = Reset(check_numerics_op.get(), "CheckNumerics",
                   /*raw_device_name=*/nullptr, &forward_op);
  ASSERT_EQ(errors::OK, s.code()) << s.message();
  if (isa<TracingOperation>(check_numerics_op.get())) {
    s = dyn_cast<TracingOperation>(check_numerics_op.get())
            ->SetOpName("check_numerics");
    ASSERT_EQ(errors::OK, s.code()) << s.message();
  }

namespace {
struct DebuggingOpts {
  // If true, insert Print nodes to print every output from an XLA cluster.
  bool print_outputs;

  // If true, insert CheckNumerics nodes for every floating point typed input to
  // an XLA cluster.
  bool check_input_numerics;

  // If true, insert CheckNumerics nodes for every floating point typed output
  // from an XLA cluster.
  bool check_output_numerics;
};

    // to cluster ControlTrigger because of how we use deadness analysis.
    bool allow_control_trigger = false;

    // Whether it is okay to "cluster" Assert and CheckNumerics by simply
    // removing them (they're not removed during clustering, but their
    // XlaOpKernel is a no-op kernel).  We avoid auto-clustering these ops so

    "BatchMatrixBandPart",
    "BatchMatrixDiag",
    "BatchMatrixDiagPart",
    "BatchMatrixSetDiag",
    "BatchToSpace",
    "BatchToSpaceND",
    "Bitcast",
    "BroadcastArgs",
    "BroadcastTo",
    "CheckNumerics",
    "ConcatV2",
    "ConjugateTranspose",
    "DebugGradientIdentity",
    "DeepCopy",
    "DepthToSpace",
    "Dequantize",
    "EditDistance",
    "Empty",
    "EnsureShape",
    "ExpandDims",

            &build_ops_flags->tf_xla_check_cluster_input_numerics,
            "If true then insert CheckNumerics nodes to check all cluster "
            "inputs."),
       Flag("tf_xla_check_cluster_output_numerics",
            &build_ops_flags->tf_xla_check_cluster_output_numerics,
            "If true then insert CheckNumerics nodes to check all cluster "
            "outputs."),
       Flag("tf_xla_disable_constant_folding",

  // clusters.  Useful for debugging.
  bool tf_xla_print_cluster_outputs;

  // If true, insert CheckNumerics nodes for every floating point typed input to
  // an XLA cluster.
  bool tf_xla_check_cluster_input_numerics;

  // If true, insert CheckNumerics nodes for every floating point typed output
  // from an XLA cluster.
  bool tf_xla_check_cluster_output_numerics;

  RunTest(x, shape, y, shape);
}

TEST_F(ArrayGradTest, CheckNumericsGrad) {
  TensorShape shape({5, 2});
  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
  auto y = CheckNumerics(scope_, x, "CheckNumerics failed");
  RunTest(x, shape, y, shape);
}

TEST_F(ArrayGradTest, ReshapeGrad) {
  TensorShape x_shape({5, 2});
  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));

def ConvertArithConstToTfConst : Pat<
  (Arith_ConstantOp:$res DenseElementsAttr:$value),
  (TF_ConstOp $value),
  [(AnyStaticShapeTensor $res)]>;

// Remove CheckNumerics op
def RemoveCheckNumerics : Pat<
  (TF_CheckNumericsOp $arg, $msg),
  (replaceWithValue $arg)>;

// Remove StopGradient op
def RemoveStopGradient : Pat<
  (TF_StopGradientOp $arg),
  (replaceWithValue $arg)>;

Example usage:

``` python
a = tf.Variable(1.0)
tf.debugging.check_numerics(a, message='')

b = tf.Variable(np.nan)
try:
  tf.debugging.check_numerics(b, message='Checking b')
except Exception as e:
  assert "Checking b : Tensor had NaN values" in e.message

c = tf.Variable(np.inf)
try:
  tf.debugging.check_numerics(c, message='Checking c')
except Exception as e:

  string err_msg = strings::StrCat(
      "Not a number (NaN) or infinity (Inf) values detected in gradient. ",
      message);
  grad_outputs->push_back(CheckNumerics(scope, grad_inputs[0], err_msg));
  return scope.status();
}
REGISTER_GRADIENT_OP("CheckNumerics", CheckNumericsGrad);

Status ReshapeGrad(const Scope& scope, const Operation& op,
                   const std::vector<Output>& grad_inputs,