Explicitly reset output parameters before the phase 2 fallback
In case, the MLIR bridge phase 2 partially succeeded, it may have populated some of the output params which should be discarded during the fallback.
PiperOrigin-RevId: 385319712
Change-Id: I930451ddd2c389d87833f109d1ae6de1fc9520ca
diff --git a/tensorflow/core/tpu/tpu_compile.cc b/tensorflow/core/tpu/tpu_compile.cc
index 2f61707..25dc2c4 100644
--- a/tensorflow/core/tpu/tpu_compile.cc
+++ b/tensorflow/core/tpu/tpu_compile.cc
@@ -252,6 +252,12 @@
std::vector<XlaCompiler::Argument>* args,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes) {
+ arg_core_mapping->clear();
+ arg_core_mapping->resize(metadata.args_size());
+
+ per_core_arg_shapes->clear();
+ per_core_arg_shapes->resize(metadata.num_cores_per_replica());
+
// Builds a description of the computation's arguments.
int constant_count = 0;
size_t guaranteed_constants_size = 0;
