output_to_input_alias[aliasing_output.getInt()] = arg_index;
  }

  if (output_to_input_alias.empty()) return absl::OkStatus();

  xla::HloModuleProto* module_proto =
      compilation_result->computation->mutable_proto();
  absl::StatusOr<xla::ProgramShape> program_shape_or_status =
      compilation_result->computation->GetProgramShape();
  TF_RET_CHECK(program_shape_or_status.ok());