From 84ff3e44b22f690c218ddf391d3f773c096d0163 Mon Sep 17 00:00:00 2001
From: Lucy Fox <lucyfox@google.com>
Date: Thu, 9 Apr 2020 21:35:39 -0700
Subject: [PATCH] Add HLO_ScalarIntTensor type to HLO dialect.

PiperOrigin-RevId: 305821689
Change-Id: Id6dc37b2050a4ae7cf334d94095fe1d6af7af559
---
 tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td b/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td
index 8dee4d0eb69..f3de67a08c1 100644
--- a/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td
+++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td
@@ -37,6 +37,11 @@ def HLO_Token : Type<CPred<"$_self.isa<TokenType>()">, "token">;
 // Any integer tensor types
 def HLO_IntTensor : TensorOf<[HLO_Int]>;
 
+// Any integer tensor type with rank 0 (i.e. representing a single integer).
+def HLO_ScalarIntTensor : ShapedContainerType<
+  [HLO_Int], And<[IsTensorTypePred, HasAnyRankOfPred<[0]>]>,
+  "a 0-dim integer tensor">;
+
 // Any floating-point tensor types
 def HLO_FpTensor : TensorOf<[AnyFloat]>;