From f4f728f18ee4a6791b32f18e454103107c1856c1 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Wed, 27 Jan 2021 05:09:01 -0800 Subject: [PATCH] Add SameOperandsAndResultShape trait to ConvertOp and IsFiniteOp. PiperOrigin-RevId: 354070005 --- include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td | 5 +++-- 1 file changed, 3 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 c69cfa5..aaa5c75 100644 --- a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td +++ b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td @@ -160,7 +160,8 @@ def HLO_CeilOp: HLO_UnaryElementwiseOp<"ceil", [NoSideEffect, SameOperandsAndResultType], HLO_FpTensor>, BASE_HLO_CeilOp; def HLO_ConvertOp : HLO_UnaryElementwiseOp<"convert", - [NoSideEffect], HLO_Tensor>, BASE_HLO_ConvertOp { + [NoSideEffect, SameOperandsAndResultShape], HLO_Tensor>, + BASE_HLO_ConvertOp { let builders = [ OpBuilderDAG<(ins "Value":$operand, "Type":$result_element_ty)>]; @@ -196,7 +197,7 @@ def HLO_ImagOp: HLO_UnaryElementwiseOp<"imag", } def HLO_IsFiniteOp: HLO_UnaryElementwiseOp<"is_finite", - [NoSideEffect], HLO_Tensor>, + [NoSideEffect, SameOperandsAndResultShape], HLO_Tensor>, BASE_HLO_IsFiniteOp { let arguments = (ins HLO_FpTensor:$x); let results = (outs HLO_PredTensor:$y);