Relax result type comparison

Check if the types are compatible insteal of equal. Sometimes the result
type can be made static and type inference cannot deduce it e.g. when
the k is dynamic

include/xten/Dialect/XTenNN/IR/XTenNNOps.td

@@ -529,6 +529,12 @@ def XtenNN_TopK: XTenNN_Op<"topk", [
   let results = (outs AnyRankedTensor:$output, AnyRankedTensor:$indices);
 
   let assemblyFormat = [{ `(`$input `:` type($input) `,` $k `:` type($k)`)` attr-dict `->` type($output) `,` type($indices) }];
+
+  let extraClassDeclaration = [{
+    /// Returns when two result types are compatible for this op; method used by
+    /// InferTypeOpInterface.
+    static bool isCompatibleReturnTypes(mlir::TypeRange l, mlir::TypeRange r);
+  }];
 }
 
 

lib/Dialect/XTenNN/IR/XTenNNOps.cpp

@@ -575,3 +575,9 @@ LogicalResult TopK::inferReturnTypeComponents(
       ShapedTypeComponents(resultShape, IntegerType::get(context, 64)));
   return success();
 }
+
+bool TopK::isCompatibleReturnTypes(mlir::TypeRange l, mlir::TypeRange r) {
+  if (l.size() != r.size() || l.size() != 2)
+    return false;
+  return succeeded(verifyCompatibleShapes(l, r));
+}