From 5b9fc81bff5086b98b60b996c1a40f3953c13fc2 Mon Sep 17 00:00:00 2001
From: Yuanzhong Xu <yuanzx@google.com>
Date: Wed, 26 Jun 2019 20:19:05 -0700
Subject: [PATCH] Add run_hlo_passes option to ExecuteReplicated without device
 assignment

PiperOrigin-RevId: 255318599
---
 tensorflow/compiler/xla/tests/hlo_test_base.cc | 3 ++-
 tensorflow/compiler/xla/tests/hlo_test_base.h  | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.cc b/tensorflow/compiler/xla/tests/hlo_test_base.cc
index 79974723b8b..64f5440b99f 100644
--- a/tensorflow/compiler/xla/tests/hlo_test_base.cc
+++ b/tensorflow/compiler/xla/tests/hlo_test_base.cc
@@ -212,9 +212,10 @@ Literal HloTestBase::ExecuteAndTransfer(std::unique_ptr<HloModule> module,
 
 StatusOr<std::vector<Literal>> HloTestBase::ExecuteReplicated(
     std::unique_ptr<HloModule> module, absl::Span<Literal* const> arguments,
-    int64 num_replicas, bool use_threads) {
+    int64 num_replicas, bool use_threads, bool run_hlo_passes) {
   HloRunner::ReplicatedExecuteOptions options;
   options.num_replicas = num_replicas;
+  options.run_hlo_passes = run_hlo_passes;
   options.use_threads = use_threads;
   for (auto argument : arguments) {
     options.arguments.push_back(argument);
diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.h b/tensorflow/compiler/xla/tests/hlo_test_base.h
index 7a78307a467..d4a1788c928 100644
--- a/tensorflow/compiler/xla/tests/hlo_test_base.h
+++ b/tensorflow/compiler/xla/tests/hlo_test_base.h
@@ -181,7 +181,7 @@ class HloTestBase : public ::testing::Test {
   // Executable::ExecuteOnStreams.
   StatusOr<std::vector<Literal>> ExecuteReplicated(
       std::unique_ptr<HloModule> module, absl::Span<Literal* const> arguments,
-      int64 num_replicas, bool use_threads);
+      int64 num_replicas, bool use_threads, bool run_hlo_passes = false);
 
   // Same as above, but uses specified device assignment.
   StatusOr<std::vector<Literal>> ExecuteReplicated(