experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add run_hlo_passes option to ExecuteReplicated without device assignment

Browse Source 
PiperOrigin-RevId: 255318599
This commit is contained in:
Yuanzhong Xu 2019-06-26 20:19:05 -07:00 committed by TensorFlower Gardener
parent 852061b75b
commit 5b9fc81bff
2 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/compiler/xla/tests/hlo_test_base.cc
View File

@ -212,9 +212,10 @@ Literal HloTestBase::ExecuteAndTransfer(std::unique_ptr<HloModule> module,
StatusOr<std::vector<Literal>> HloTestBase::ExecuteReplicated( StatusOr<std::vector<Literal>> HloTestBase::ExecuteReplicated(
std::unique_ptr<HloModule> module, absl::Span<Literal* const> arguments, 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; HloRunner::ReplicatedExecuteOptions options;
options.num_replicas = num_replicas; options.num_replicas = num_replicas;
options.run_hlo_passes = run_hlo_passes;
options.use_threads = use_threads; options.use_threads = use_threads;
for (auto argument : arguments) { for (auto argument : arguments) {
options.arguments.push_back(argument); options.arguments.push_back(argument);

2
tensorflow/compiler/xla/tests/hlo_test_base.h
View File

@ -181,7 +181,7 @@ class HloTestBase : public ::testing::Test {
// Executable::ExecuteOnStreams. // Executable::ExecuteOnStreams.
StatusOr<std::vector<Literal>> ExecuteReplicated( StatusOr<std::vector<Literal>> ExecuteReplicated(
std::unique_ptr<HloModule> module, absl::Span<Literal* const> arguments, 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. // Same as above, but uses specified device assignment.
StatusOr<std::vector<Literal>> ExecuteReplicated( StatusOr<std::vector<Literal>> ExecuteReplicated(