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

[XLA:Python] Expose new use_spmd_partitioning compile option.

Browse Source 
PiperOrigin-RevId: 311761373
Change-Id: I1f696e0c082295dc0e6896f05d1e88525de7ce70
This commit is contained in:
Skye Wanderman-Milne 2020-05-15 10:50:40 -07:00 committed by TensorFlower Gardener
parent 262e92804b
commit 26104505b8
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/compiler/xla/python/xla.cc
View File

@ -1406,7 +1406,10 @@ PYBIND11_MODULE(xla_extension, m) {
options.device_assignment())
: absl::nullopt;
},
&ExecutableBuildOptions::set_device_assignment);
&ExecutableBuildOptions::set_device_assignment)
.def_property("use_spmd_partitioning",
&ExecutableBuildOptions::use_spmd_partitioning,
&ExecutableBuildOptions::set_use_spmd_partitioning);
py::class_<XlaComputation>(m, "XlaComputation")
.def(py::init([](const py::bytes& serialized_hlo_module_proto)