Minor fixes

include/torch-mlir/Dialect/Torch/IR/GeneratedPrimOps.td

@@ -175,6 +175,7 @@ def Torch_PrimRaiseExceptionOp : Torch_Op<"prim.RaiseException", [
 }
 
 def Torch_PrimUninitializedOp : Torch_Op<"prim.Uninitialized", [
+    NoSideEffect,
     AllowsTypeRefinement,
     HasValueSemantics
   ]> {
@@ -185,7 +186,6 @@ def Torch_PrimUninitializedOp : Torch_Op<"prim.Uninitialized", [
     AnyTorchType:$result
   );
   let assemblyFormat = " attr-dict `:` qualified(type($result))";
-  let hasCanonicalizer = 1;
 }
 
 def Torch_PrimUncheckedCastOp : Torch_Op<"prim.unchecked_cast", [

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -1037,20 +1037,6 @@ void PrimTupleIndexOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
   });
 }
 
-//===----------------------------------------------------------------------===//
-// PrimUninitializedOp
-//===----------------------------------------------------------------------===//
-
-void PrimUninitializedOp::getCanonicalizationPatterns(
-    RewritePatternSet &patterns, MLIRContext *context) {
-  patterns.add(+[](PrimUninitializedOp op, PatternRewriter &rewriter) {
-    if (!op.use_empty())
-      return failure();
-    rewriter.eraseOp(op);
-    return success();
-  });
-}
-
 //===----------------------------------------------------------------------===//
 // PrimTupleUnpackOp
 //===----------------------------------------------------------------------===//

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -414,7 +414,7 @@ def emit_prim_ops(torch_ir_dir: str, registry: Registry):
         emit("prim::max.self_int : (int[]) -> (int)")
         emit("prim::max.int : (int, int) -> (int)")
         emit("prim::RaiseException : (str) -> ()")
-        emit("prim::Uninitialized : () -> (Any)", has_canonicalizer=True)
+        emit("prim::Uninitialized : () -> (Any)", traits=["NoSideEffect"])
         emit("prim::unchecked_cast : (t) -> (t)",
              traits=["DeclareOpInterfaceMethods<CastOpInterface>"])
         emit("prim::Print : (...) -> ()")

python/torch_mlir/dialects/torch/importer/jit_ir/csrc/torch_to_mlir_utils.cpp

@@ -11,6 +11,7 @@
 #include "function_importer.h"
 #include "ivalue_importer.h"
 
+#include <ATen/TensorUtils.h>
 #include <unordered_map>
 
 #include "mlir_utils.h"