From b2418b7feeeb0530a7c955eb61eaf8354ca432a8 Mon Sep 17 00:00:00 2001 From: Sean Silva Date: Tue, 14 Apr 2020 01:00:05 -0700 Subject: [PATCH] Allow index inside dimension tensors. Ever since https://reviews.llvm.org/rG3dceb6d2460550ef8c2e53df2181c8a11a7e6730 MLIR now allows `index` type inside tensors, addressing this longstanding TODO. PiperOrigin-RevId: 306391413 Change-Id: I797623de02ff031b9ffabfaee81afba91ead6692 --- tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td b/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td index f3de67a08c1..7994026ac3b 100644 --- a/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td +++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td @@ -55,11 +55,10 @@ def HLO_Tuple : NestedTupleOf<[HLO_Tensor, HLO_Token]>; def HLO_TensorOrTuple : AnyTypeOf<[HLO_Tensor, HLO_Tuple]>; -// Dynamic representation of a shape vector as a tensor. Ideally this would be -// an index type (as it stores indices) but that is currently disallowed in -// MLIR. +// Dynamic representation of a shape vector as a tensor. def HLO_DimensionTensor : ShapedContainerType< - [AnySignlessInteger], And<[IsTensorTypePred, HasAnyRankOfPred<[1]>]>, + [Index, AnySignlessInteger], + And<[IsTensorTypePred, HasAnyRankOfPred<[1]>]>, "a 1D tensor of dimensions">; // In general, static shaped tensor constraints should be avoided unless