diff options
author | 2018-06-08 12:50:16 -0700 | |
---|---|---|
committer | 2018-06-08 12:52:42 -0700 | |
commit | 278fbe4146b160980fec318187546d9d8870d244 (patch) | |
tree | 402e929ee51c089458fa7bfac1ae6826f73d71cd /tensorflow/compiler/xla/service/shape_inference.cc | |
parent | 7bb79ee219d4efbd92d1ef4e0dbe45f4aee26654 (diff) |
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
Diffstat (limited to 'tensorflow/compiler/xla/service/shape_inference.cc')
-rw-r--r-- | tensorflow/compiler/xla/service/shape_inference.cc | 11 |
1 files changed, 11 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/shape_inference.cc b/tensorflow/compiler/xla/service/shape_inference.cc index d624f548b1..fdc7f41759 100644 --- a/tensorflow/compiler/xla/service/shape_inference.cc +++ b/tensorflow/compiler/xla/service/shape_inference.cc @@ -463,6 +463,17 @@ StatusOr<Shape> InferWindowOutputShape(const Shape& base_shape, return ShapeUtil::MakeShape(element_type, new_dimensions); } +/* static */ StatusOr<Shape> ShapeInference::InferTokenShape( + tensorflow::gtl::ArraySlice<const Shape*> arg_shapes) { + for (const Shape* arg_shape : arg_shapes) { + if (arg_shape->element_type() != TOKEN) { + return InvalidArgument( + "Operands of token instructions must be TOKEN types."); + } + } + return ShapeUtil::MakeTokenShape(); +} + /* static */ StatusOr<Shape> ShapeInference::InferConvertShape( const Shape& operand_shape, PrimitiveType new_element_type) { auto old_element_type = operand_shape.element_type(); |