[XLA] Fix the parameter instruction printing issue

Append the parameter number to the fusion parameter name, and use the parameter
name rather the instruction name in creating the new parameter.

Show the paramameter number when printing out parameter instructions.
Change: 153752424

1 files changed, 7 insertions, 3 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc
index 526c1c58b3..0fb8c06f88 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.cc
+++ b/tensorflow/compiler/xla/service/hlo_computation.cc
@@ -93,6 +93,10 @@ HloInstruction* HloComputation::AddInstructionInternal(
   // Generate a unique name for the instruction.
   instruction->set_name(
       instruction_name_uniquer_.GetUniqueName(instruction->name()));
+  if (instruction->opcode() == HloOpcode::kParameter) {
+    instruction->set_parameter_name(
+        instruction_name_uniquer_.GetUniqueName(instruction->parameter_name()));
+  }
   Reparent(instruction.get());
   HloInstruction* pinst = instruction.get();
   instruction_iterators_[pinst] =
@@ -127,9 +131,9 @@ Status HloComputation::RemoveParameter(int64 param_no) {
 
   while (param_no < param_instructions_.size()) {
     param_instruction = param_instructions_[param_no];
-    HloInstruction* new_instr =
-        AddInstructionInternal(HloInstruction::CreateParameter(
-            param_no, param_instruction->shape(), param_instruction->name()));
+    HloInstruction* new_instr = AddInstructionInternal(
+        HloInstruction::CreateParameter(param_no, param_instruction->shape(),
+                                        param_instruction->parameter_name()));
     TF_RETURN_IF_ERROR(param_instruction->ReplaceAllUsesWith(new_instr));
     new_instr->SetParentFusion(root_instruction_->fusion_instruction());
     param_instructions_[param_no] = new_instr;