message: [XLA] Propagate the memory space to tuple outputs.
diff: @@ -254,10 +254,12 @@ void BufferAllocation::AddAssignment(const HloValue& buffer, int64 offset,
  assigned_buffers_.emplace(&buffer, offset_size);
  // For debugging purposes, store the assigned memory space in the
  // instruction's layout.
  - HloInstruction* defining_instruction = buffer.defining_instruction();
  - if (defining_instruction->shape().has_layout()) {
  - defining_instruction->mutable_shape()->mutable_layout()->set_memory_space(
  - buffer.color().value());
  + for (HloPosition position : buffer.positions()) {
  + Shape* shape = ShapeUtil::GetMutableSubshape(
  + position.instruction->mutable_shape(), position.index);
  + if (shape->has_layout()) {
  + shape->mutable_layout()->set_memory_space(buffer.color().value());
  + }
  }
  }
  