Remove assertions checking device map size in TPUVariableRuntimeReformattingPass.
It is possible for more than one replicated device to be populated (e.g. replicated host).
PiperOrigin-RevId: 317419924
Change-Id: I1d8fa95654f324557d78adf4570adc6c3cfdabb4
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_variable_runtime_reformatting.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_variable_runtime_reformatting.cc
index ec4a25c..d88982d 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_variable_runtime_reformatting.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_variable_runtime_reformatting.cc
@@ -261,7 +261,6 @@
// placed in logical core 0.
// TODO(b/148913020): Remove this constraint once model parallelism is
// supported.
- assert(devices.size() == 1);
assert(devices.find(tensorflow::GetDeviceAliasForLogicalCore(0))
->getSecond()
.size() == num_replicas);
@@ -369,9 +368,6 @@
// TODO(b/148913020): Remove this constraint once model parallelism is
// supported.
- assert(devices.size() == 1 &&
- "As model parallelism is not supported yet, tf_device.replicate "
- "`devices` attribute should have one dictionary element.");
const auto& device_list =
devices.find(tensorflow::GetDeviceAliasForLogicalCore(0))->getSecond();
