From b97e5034da8c509fab11e848ce4bbd489f7e0b59 Mon Sep 17 00:00:00 2001
From: Henry Tu <henry.tu@cerebras.net>
Date: Thu, 25 Aug 2022 09:29:45 -0400
Subject: [PATCH 1/2] Revert updating mlir_native_functions.cpp signature, due
 to
 https://github.com/pytorch/pytorch/commit/a7edf713608806f10e17eab90d0a5df727d9a16e

---
 .../mlir_native_functions.cpp                 | 29 +++++--------------
 1 file changed, 7 insertions(+), 22 deletions(-)

diff --git a/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp b/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp
index cfd0cf68e15a..ec72c37ab508 100644
--- a/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp
+++ b/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp
@@ -302,14 +302,10 @@ at::Tensor LazyNativeFunctions::_to_copy(
 };
 
 at::Tensor LazyNativeFunctions::empty(
-    at::SymIntArrayRef sym_size,
-    c10::optional<at::ScalarType> dtype,
-    c10::optional<at::Layout> layout,
-    c10::optional<at::Device> device,
+    at::IntArrayRef size, c10::optional<at::ScalarType> dtype,
+    c10::optional<at::Layout> layout, c10::optional<at::Device> device,
     c10::optional<bool> pin_memory,
     c10::optional<at::MemoryFormat> memory_format) {
-  // TODO: support this directly
-  auto size = c10::asIntArrayRefSlow(sym_size);
   const auto device_type = torch::lazy::getBackend()->EagerFallbackDeviceType();
   at::TensorOptions options = at::TensorOptions()
                                   .device(c10::Device(device_type))
@@ -321,9 +317,8 @@ at::Tensor LazyNativeFunctions::empty(
   // See Note [Lazy Tensor Functionalization]
   if (c10::impl::tls_local_dispatch_key_set().excluded_.has(
           c10::DispatchKey::Functionalize)) {
-    // Invariant: if the functionalization key is in the exclude set, then we're
-    // expected to return an ordinary tensor, which will be "lifted" into a
-    // functional wrapper later.
+    // Invariant: if the functionalization key is in the exclude set, then we're expected
+    // to return an ordinary tensor, which will be "lifted" into a functional wrapper later.
     return tensor;
   } else {
     auto wrapped = at::functionalization::impl::to_functional_tensor(tensor);
@@ -336,13 +331,7 @@ at::Tensor LazyNativeFunctions::empty_strided(
     c10::optional<at::ScalarType> dtype, c10::optional<at::Layout> layout,
     c10::optional<at::Device> device, c10::optional<bool> pin_memory) {
   TORCH_LAZY_FN_COUNTER("lazy::");
-  at::Tensor t = empty(
-      c10::SymIntArrayRef::fromIntArrayRef(size),
-      dtype,
-      layout,
-      device,
-      pin_memory,
-      c10::nullopt);
+  at::Tensor t = empty(size, dtype, layout, device, pin_memory, c10::nullopt);
   return t.as_strided(size, stride, /*storage_offset=*/0);
 }
 
@@ -361,8 +350,7 @@ LazyNativeFunctions::fill_(at::Tensor& self, const at::Scalar& value) {
 at::Tensor LazyNativeFunctions::_unsafe_view(
     const at::Tensor& self, at::IntArrayRef size) {
   TORCH_LAZY_FN_COUNTER("lazy::");
-  return LazyNativeFunctions::view_copy(
-      self, c10::SymIntArrayRef::fromIntArrayRef(size));
+  return LazyNativeFunctions::view_copy(self, size);
 }
 
 // This is needed by the torch.tensor constructor.
@@ -398,10 +386,7 @@ at::Tensor LazyNativeFunctions::new_empty_strided(
 }
 
 at::Tensor LazyNativeFunctions::narrow_copy(
-    const at::Tensor& self,
-    int64_t dim,
-    c10::SymInt start,
-    c10::SymInt length) {
+    const at::Tensor& self, int64_t dim, int64_t start, int64_t length) {
   return at::functionalization::functionalize_aten_op<ATEN_OP(
       narrow_copy)>::call(self, dim, start, length);
 }

From 879d010a718e3e05de1b6267b5fff3b68425b2a0 Mon Sep 17 00:00:00 2001
From: Henry Tu <henry.tu@cerebras.net>
Date: Thu, 25 Aug 2022 10:00:17 -0400
Subject: [PATCH 2/2] Restored NewZeros to LTC XFAIL set

---
 e2e_testing/torchscript/xfail_sets.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/e2e_testing/torchscript/xfail_sets.py b/e2e_testing/torchscript/xfail_sets.py
index f106c974910a..8904d7faa91e 100644
--- a/e2e_testing/torchscript/xfail_sets.py
+++ b/e2e_testing/torchscript/xfail_sets.py
@@ -435,6 +435,12 @@
     "NewOnesModuleFloat3D_basic",
     "NewOnesModuleInt2D_basic",
     "NewOnesModuleInt3D_basic",
+    "NewZerosModuleDefaultDtype_basic",
+    "NewZerosModuleFalsePinMemory_basic",
+    "NewZerosModuleFloat2D_basic",
+    "NewZerosModuleFloat3D_basic",
+    "NewZerosModuleInt2D_basic",
+    "NewZerosModuleInt3D_basic",
     "OnesLikeModule_defaultDtype",
     "OnesLikeModule_falsePinMemory",
     "OnesLikeModule_float",