From de5ddaf7c984fc03712098c30d50bcd435fe3d77 Mon Sep 17 00:00:00 2001
From: Robert Suderman <suderman@google.com>
Date: Wed, 29 Jul 2020 12:31:15 -0700
Subject: [PATCH] HLO Random operations should match shape constraints.

PiperOrigin-RevId: 323844002
---
 include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
index 0ed4235..93c5388 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
+++ b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
@@ -1324,7 +1324,7 @@ def HLO_RngUniformOp : HLO_Op<"rng_uniform", []>, BASE_HLO_RngUniformOp {
   let arguments = (ins
     HLO_PredIntOrFpTensor:$a,
     HLO_PredIntOrFpTensor:$b,
-    I64Tensor:$shape
+    HLO_DimensionTensor:$shape
   );
 
   let results = (outs HLO_PredIntOrFpTensor);
@@ -1336,7 +1336,7 @@ def HLO_RngNormalOp : HLO_Op<"rng_normal", []>, BASE_HLO_RngNormalOp {
   let arguments = (ins
     HLO_FpTensor:$mu,
     HLO_FpTensor:$sigma,
-    I64Tensor:$shape
+    HLO_DimensionTensor:$shape
   );
 
   let results = (outs HLO_FpTensor);