use verifyAtLeastNOperands trait

openxla · Feb 29, 2024 · 2a8fcd4 · 2a8fcd4
1 parent 2b01730
commit 2a8fcd4
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/stablehlo/dialect/Base.h b/stablehlo/dialect/Base.h
@@ -312,10 +312,10 @@ class CompatibleOperandsElementType
                                       CompatibleOperandsElementType> {
  public:
   static LogicalResult verifyTrait(Operation *op) {
-    Type expected;
-    if (op->getNumOperands() != 0) expected = op->getOperand(0).getType();
-    if (!expected) return failure();
+    if (failed(mlir::OpTrait::impl::verifyAtLeastNOperands(op, 1)))
+      return failure();
 
+    Type expected = op->getOperand(0).getType();
     auto typeMatch = [&](Type actual) {
       return isCompatibleElementTypeForHloTypeInference(actual, expected);
     };