Add kGenerateToken HLO instruction.

The new HLO instruction serves two purposes. (1) It generates a new token value. This is the only way to create tokens. (2) The operation is variadic, taking zero or more token operands. The operation acts as a join of its operands.

I considered initially using a kConstant constant as a method to create new tokens, but this ran into problems because of expectations in backends regarding constants and their materialization.

This CL enables creation of generate-token instructions, but the new instruction is not supported yet in any backend.

PiperOrigin-RevId: 199836205

1 files changed, 8 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.cc b/tensorflow/compiler/xla/service/hlo_evaluator.cc
index 1e78d775c8..e0648e1467 100644
--- a/tensorflow/compiler/xla/service/hlo_evaluator.cc
+++ b/tensorflow/compiler/xla/service/hlo_evaluator.cc
@@ -910,6 +910,14 @@ Status HloEvaluator::HandleBroadcast(HloInstruction* broadcast) {
   return Status::OK();
 }
 
+Status HloEvaluator::HandleGenerateToken(HloInstruction* token) {
+  // Literals cannot represent a TOKEN shape so just create an empty tuple as
+  // the "result" of the kGenerateToken operation.
+  // TODO(b/109929053): Add support for TOKENs in Literals.
+  evaluated_[token] = Literal::MakeTuple({});
+  return Status::OK();
+}
+
 Status HloEvaluator::HandleGetTupleElement(HloInstruction* get_tuple_element) {
   const auto result_shape = get_tuple_element->shape();
   const int64 index = get_tuple_element->tuple_index();