[XLA] Add and use a layout-sensitive HLO verifier.

For now, this verifier checks some noncontroversial invariants, like:

 - Fusion operands and fusion computation parameters must have matching layouts.
 - Same for while loops, calls, kConditional.

It's a bit of a pain to add these explicit layout-sensitive and
allow-mixed-precision flags everywhere, but I think it's better than
adding default args.  With default args we can easily mix up the order,
and we'd only be able to add new flags to the end of the list.

PiperOrigin-RevId: 210059349

1 files changed, 6 insertions, 2 deletions

diff --git a/tensorflow/compiler/xla/tests/test_utils.cc b/tensorflow/compiler/xla/tests/test_utils.cc
index 2f1d97b25d..21c58e075e 100644
--- a/tensorflow/compiler/xla/tests/test_utils.cc
+++ b/tensorflow/compiler/xla/tests/test_utils.cc
@@ -408,8 +408,12 @@ StatusOr<std::vector<std::unique_ptr<Literal>>> MakeFakeArguments(
   return std::move(arguments);
 }
 
-Status VerifyHloModule(HloModule* const module, bool allow_mixed_precision) {
-  return HloVerifier(allow_mixed_precision).Run(module).status();
+Status VerifyHloModule(HloModule* const module, bool layout_sensitive,
+                       bool allow_mixed_precision) {
+  return HloVerifier(/*layout_sensitive=*/layout_sensitive,
+                     /*allow_mixed_precision=*/allow_mixed_precision)
+      .Run(module)
+      .status();
 }
 
 }  // namespace xla