[tfdbg] Add debugging.enable_check_numerics() to mnist example

When running the mnist example with --debug, it will halt with an InvalidArgumentError and a message which shows the stack trace which called the operation that outputed an Inf/NaN.

PiperOrigin-RevId: 269815297

tensorflow/python/debug/BUILD

@@ -1263,7 +1263,7 @@ sh_test(
 
 sh_test(
     name = "examples_v2_test",
-    size = "medium",
+    size = "small",
     srcs = ["examples/v2/examples_v2_test.sh"],
     data = [
         ":debug_fibonacci_v2",

tensorflow/python/debug/examples/v2/debug_fibonacci_v2.py

@@ -42,7 +42,9 @@ def main(_):
         "tfdbg v2 support for debug_fibonacci is not implemented yet")
   elif FLAGS.tensorboard_debug_address:
     raise NotImplementedError(
-        "tfdbg v2 support for debug_fibonacci is not implemented yet")
+        "Tensorboard Debugger Plugin support for debug_fibonacci_v2 is not "
+        "implemented yet"
+    )
 
   # Construct the TensorFlow network.
   n0 = tf.constant(np.ones([FLAGS.tensor_size] * 2), dtype=tf.int32)

tensorflow/python/debug/examples/v2/debug_mnist_v2.py

@@ -111,16 +111,17 @@ def main(_):
         "The --debug and --tensorboard_debug_address flags are mutually "
         "exclusive.")
   if FLAGS.debug:
-    raise NotImplementedError(
-        "tfdbg v2 support for debug_mnist is not implemented yet")
+    tf.debugging.enable_check_numerics()
   elif FLAGS.tensorboard_debug_address:
     raise NotImplementedError(
-        "tfdbg v2 support for debug_mnist is not implemented yet")
+        "Tensorboard Debugger Plugin support for debug_mnist_v2 is not "
+        "implemented yet"
+    )
 
   # Import data
   if FLAGS.fake_data:
-    imgs = tf.random.uniform(maxval=256, shape=(10, 28, 28), dtype=tf.int32)
-    labels = tf.random.uniform(maxval=10, shape=(10,), dtype=tf.int32)
+    imgs = tf.random.uniform(maxval=256, shape=(1000, 28, 28), dtype=tf.int32)
+    labels = tf.random.uniform(maxval=10, shape=(1000,), dtype=tf.int32)
     mnist_train = imgs, labels
     mnist_test = imgs, labels
   else:

tensorflow/python/debug/examples/v2/examples_v2_test.sh

@@ -51,18 +51,17 @@ else
   DEBUG_MNIST_BIN="${PYTHON_BIN_PATH} -m tensorflow.python.debug.examples.v2.debug_mnist"
 fi
 
-# Override the default ui_type=curses to allow the test to pass in a tty-less
-# test environment.
-cat << EOF | ${DEBUG_FIBONACCI_BIN} --tensor_size=2
-run
-exit
-EOF
+# Verify fibonacci runs normally without additional flags
+${DEBUG_FIBONACCI_BIN} --tensor_size=2
 
-cat << EOF | ${DEBUG_MNIST_BIN} --max_steps=1 --fake_data
-run -t 1
-run --node_name_filter hidden --op_type_filter MatMul
-run -f has_inf_or_nan
-EOF
+# Verify mnist runs normally without additional flags
+${DEBUG_MNIST_BIN} --max_steps=4 --fake_data
+
+# Verify mnist does not break with check_numerics enabled on first iteration
+# check_numerics should not cause non-zero exit code on a single train step
+${DEBUG_MNIST_BIN} --max_steps=1 --fake_data --debug
+
+# Verify check_numerics exits with non-zero exit code
+! ${DEBUG_MNIST_BIN} --max_steps=4 --fake_data --debug
 
-echo
 echo "SUCCESS: tfdbg examples and binaries test PASSED"