From 3dc97120136c5c9ca332c7d56878187281bc6233 Mon Sep 17 00:00:00 2001
From: Smit Hinsu <hinsu@google.com>
Date: Wed, 29 Apr 2020 17:47:54 -0700
Subject: [PATCH] Allow unsigned integers in HLO_DimensionTensor and
 HLO_StaticShapeTensor

Follow up after enabling unsigned integers in the dialect.

PiperOrigin-RevId: 309134014
Change-Id: Ie458c93ec69a5ed1be183fdd208dc27d3de61e1a
---
 tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td b/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td
index 44e0abab031..c087ffd1f40 100644
--- a/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td
+++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td
@@ -64,14 +64,14 @@ def HLO_TensorOrTuple : AnyTypeOf<[HLO_Tensor, HLO_Tuple]>;
 
 // Dynamic representation of a shape vector as a tensor.
 def HLO_DimensionTensor : ShapedContainerType<
-    [Index, AnySignlessInteger],
+    [Index, HLO_Pred, HLO_Int],
     And<[IsTensorTypePred, HasAnyRankOfPred<[1]>]>,
     "a 1D tensor of dimensions">;
 
 // In general, static shaped tensor constraints should be avoided unless
 // it is for a legacy op which is only correct with static shapes.
 def HLO_StaticShapeTensor : StaticShapeTensorOf<[
-      AnyFloat, AnySignlessInteger, HLO_Complex]>;
+      AnyFloat, HLO_Pred, HLO_Int, HLO_Complex]>;
 
 //===----------------------------------------------------------------------===//
 // XLA on tensors combined type definitions.