Add SameOperandsAndResultShape trait to ConvertOp and IsFiniteOp.
PiperOrigin-RevId: 354070005
Change-Id: Ie1f50c51904b1e5e099c7be299dc9a04562e6ec5
diff --git a/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td b/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
index c69cfa5..aaa5c75 100644
--- a/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
+++ b/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
@@ -160,7 +160,8 @@
[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_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);
