Allow TPU's special FP debug flags to be passed via DebugOptions,
PiperOrigin-RevId: 304934290
Change-Id: I628236b106a33f651c2ba7b2b8a1f93432291cc0
diff --git a/tensorflow/compiler/xla/debug_options_flags.cc b/tensorflow/compiler/xla/debug_options_flags.cc
index 9808dd3..8604531 100644
--- a/tensorflow/compiler/xla/debug_options_flags.cc
+++ b/tensorflow/compiler/xla/debug_options_flags.cc
@@ -535,6 +535,16 @@
bool_setter_for(&DebugOptions::set_xla_gpu_deterministic_reductions),
flag_values->xla_gpu_deterministic_reductions(),
"Always run deterministic reductions on GPU"),
+ tensorflow::Flag(
+ "xla_tpu_detect_nan",
+ bool_setter_for(&DebugOptions::set_xla_tpu_detect_nan),
+ flag_values->xla_tpu_detect_nan(),
+ "Trigger error on execution on TPU if a NAN value is detected"),
+ tensorflow::Flag(
+ "xla_tpu_detect_inf",
+ bool_setter_for(&DebugOptions::set_xla_tpu_detect_inf),
+ flag_values->xla_tpu_detect_inf(),
+ "Trigger error on execution on TPU if a INF value is detected"),
});
ParseFlagsFromEnvAndDieIfUnknown("XLA_FLAGS", *flag_objects);
}
diff --git a/tensorflow/compiler/xla/xla.proto b/tensorflow/compiler/xla/xla.proto
index 3468c12..f8bd7a0 100644
--- a/tensorflow/compiler/xla/xla.proto
+++ b/tensorflow/compiler/xla/xla.proto
@@ -265,7 +265,11 @@
// Guarantee run-to-run determinism from reductions on XLA:GPU.
bool xla_gpu_deterministic_reductions = 130;
- // Next id: 135
+ // Debug options that trigger execution errors when NaN or Inf are detected.
+ bool xla_tpu_detect_nan = 135;
+ bool xla_tpu_detect_inf = 136;
+
+ // Next id: 137
// Extra options to pass to the compilation backend (e.g. LLVM); specific
// interpretation of these values is left to the backend.
