Fix failing test

Xilinx · Aug 21, 2024 · e995188 · e995188
1 parent 81503ea
commit e995188
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/lib/Dialect/XTenNN/IR/XTenNNOps.cpp b/lib/Dialect/XTenNN/IR/XTenNNOps.cpp
@@ -579,5 +579,9 @@ LogicalResult TopK::inferReturnTypeComponents(
 bool TopK::isCompatibleReturnTypes(mlir::TypeRange l, mlir::TypeRange r) {
   if (l.size() != r.size() || l.size() != 2)
     return false;
-  return succeeded(verifyCompatibleShapes(l, r));
+
+  auto sameElementType =
+      getElementTypeOrSelf(l[0]) == getElementTypeOrSelf(r[0]) &&
+      getElementTypeOrSelf(l[1]) == getElementTypeOrSelf(r[1]);
+  return sameElementType && succeeded(verifyCompatibleShapes(l, r));
 }