Allow element type traits to operate on scalars

This allows confirming that a scalar argument has the same element type as a shaped one. It's easy to validate a type is shaped on its own if that's desirable, so this shouldn't make that use case harder. This matches the behavior of other traits that operate on element type (e.g. AllElementTypesMatch). Also this makes the code simpler because now we just use getElementTypeOrSelf.

Verified that all uses in core already check the type is shaped in another way.

PiperOrigin-RevId: 273068507

1 files changed, 4 insertions, 4 deletions

diff --git a/mlir/test/lib/TestDialect/TestOps.td b/mlir/test/lib/TestDialect/TestOps.td
index a7c5fa05235..dd620def6b0 100644
--- a/mlir/test/lib/TestDialect/TestOps.td
+++ b/mlir/test/lib/TestDialect/TestOps.td
@@ -236,14 +236,14 @@ def FunctionalRegionOp : TEST_Op<"functional_region_op",
 
 def SameOperandElementTypeOp : TEST_Op<"same_operand_element_type",
     [SameOperandsElementType]> {
-  let arguments = (ins AnyVectorOrTensor, AnyVectorOrTensor);
-  let results = (outs AnyVectorOrTensor);
+  let arguments = (ins AnyType, AnyType);
+  let results = (outs AnyType);
 }
 
 def SameOperandAndResultElementTypeOp : TEST_Op<"same_operand_and_result_element_type",
     [SameOperandsAndResultElementType]> {
-  let arguments = (ins Variadic<AnyVectorOrTensor>);
-  let results = (outs Variadic<AnyVectorOrTensor>);
+  let arguments = (ins Variadic<AnyType>);
+  let results = (outs Variadic<AnyType>);
 }
 
 def SameOperandShapeOp : TEST_Op<"same_operand_shape", [SameOperandsShape]> {