[XLA] Add HLO matchers that check parameter numbers and GTE indices.

This lets you do

  EXPECT_THAT(foo, op::Parameter(42));

and

  EXPECT_THAT(bar, op::GetTupleElement(baz, 8));

PiperOrigin-RevId: 174113597

1 files changed, 29 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_matchers.cc b/tensorflow/compiler/xla/service/hlo_matchers.cc
index 0660d5a182..4255d60866 100644
--- a/tensorflow/compiler/xla/service/hlo_matchers.cc
+++ b/tensorflow/compiler/xla/service/hlo_matchers.cc
@@ -73,6 +73,35 @@ void HloMatcher::DescribeTo(::std::ostream* os) const {
   }
 }
 
+bool HloParameterMatcher::MatchAndExplain(
+    const HloInstruction* instruction,
+    ::testing::MatchResultListener* listener) const {
+  if (!HloMatcher::MatchAndExplain(instruction, listener)) {
+    return false;
+  }
+  if (instruction->parameter_number() != parameter_number_) {
+    *listener << "has wrong parameter number (got "
+              << instruction->parameter_number() << ", want "
+              << parameter_number_ << ")";
+    return false;
+  }
+  return true;
+}
+
+bool HloGetTupleElementMatcher::MatchAndExplain(
+    const HloInstruction* instruction,
+    ::testing::MatchResultListener* listener) const {
+  if (!HloMatcher::MatchAndExplain(instruction, listener)) {
+    return false;
+  }
+  if (instruction->tuple_index() != tuple_index_) {
+    *listener << "has wrong tuple index (got " << instruction->tuple_index()
+              << ", want " << tuple_index_ << ")";
+    return false;
+  }
+  return true;
+}
+
 }  // namespace testing
 
 void PrintTo(const HloInstruction* inst, ::std::ostream* os) {