move torch._functionalize APIs to pybind. add one for marking storage…

… mutations (pytorch#132337) Pull Request resolved: pytorch#132337 Approved by: https://github.com/albanD, https://github.com/justinchuby ghstack dependencies: pytorch#132243
kashif · Aug 5, 2024 · 1a0db29 · 1a0db29
1 parent 4db368a
commit 1a0db29
Show file tree

Hide file tree

Showing 3 changed files with 177 additions and 554 deletions.
diff --git a/aten/src/ATen/FunctionalTensorWrapper.h b/aten/src/ATen/FunctionalTensorWrapper.h
@@ -161,6 +161,10 @@ struct TORCH_API FunctionalTensorWrapper : public c10::TensorImpl {
     return was_storage_changed_;
   }
 
+  void set_storage_changed() {
+    was_storage_changed_ = true;
+  }
+
   c10::SymInt get_storage_size(bool before) {
     return functional_storage_impl()->get_storage_size(before);
   }

diff --git a/test/onnx/test_fx_to_onnx_with_onnxruntime.py b/test/onnx/test_fx_to_onnx_with_onnxruntime.py
@@ -539,7 +539,7 @@ def forward(self, x):
         )
 
     @pytorch_test_common.xfail_if_model_type_is_not_exportedprogram(
-        error_message="at::functionalization::impl::isFunctionalTensor(self_) INTERNAL ASSERT FAILED"
+        error_message="at::functionalization::impl::isFunctionalTensor(t) INTERNAL ASSERT FAILED"
     )
     def test_expand_as_fill_separate_tensor(self):
         class Model(torch.nn.Module):