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, 7 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/shape_inference.h b/tensorflow/compiler/xla/service/shape_inference.h
index 9da2c99b41..6100e2cd33 100644
--- a/tensorflow/compiler/xla/service/shape_inference.h
+++ b/tensorflow/compiler/xla/service/shape_inference.h
@@ -227,6 +227,13 @@ class ShapeInference {
   static StatusOr<Shape> InferConcatOpShape(
       tensorflow::gtl::ArraySlice<const Shape*> arg_shapes, int64 dimension);
 
+  // Infers the shape produced by a kGenerateToken operation. Trivially this
+  // shape is always a TOKEN shape. However, ShapeInference serves two purposes:
+  // inferring shapes and checking operand shapes. This method verifies that the
+  // operand shapes are all TOKENs.
+  static StatusOr<Shape> InferTokenShape(
+      tensorflow::gtl::ArraySlice<const Shape*> arg_shapes);
+
   // Helper that validates the given operand shape can be converted to the
   // target output_shape via a convert instruction -- the requirement is that
   // the shape is identical except for the element type.