[retake] Use extension module for pybind (#14029)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #14024 by
@lucylq
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/lucylq/106/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/lucylq/106/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/lucylq/106/orig
@diff-train-skip-merge

---------

Co-authored-by: lucylq <[email protected]>

Author: pytorchbot

exir/backend/test/test_backends.py

@@ -266,7 +266,6 @@ def forward(self, a, x, b):
 
         executorch_module = _load_for_executorch_from_buffer(buff)
 
-        # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`.
         inputs_flattened, _ = tree_flatten(model_inputs)
         model_output = executorch_module.run_method("forward", tuple(inputs_flattened))
         ref_output = add_mul_module(*model_inputs)
@@ -319,16 +318,17 @@ def forward(self, x):
         )
 
         buff = exec_prog.buffer
-
+        executorch_module = _load_for_executorch_from_buffer(buff)
         # This line should raise an exception like
         # RuntimeError: failed with error 0x12
-        _load_for_executorch_from_buffer(buff)
+        inputs_flattened, _ = tree_flatten(model_inputs)
+        executorch_module.run_method("forward", tuple(inputs_flattened))
 
     @vary_segments
     def test_backend_with_compiler_out_of_range(self, extract_delegate_segments: bool):
         with self.assertRaisesRegex(
             RuntimeError,
-            "loading method forward failed with error 0x12",
+            "Failed to execute method forward, error: 0x12",
         ):
             self.run_model_in_unsupported_backend(
                 extract_delegate_segments=extract_delegate_segments
@@ -876,7 +876,6 @@ def forward(self, a, x, b):
         self.assertEqual(counter, 2)
 
         executorch_module = _load_for_executorch_from_buffer(executorch_prog.buffer)
-        # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`.
         inputs_flattened, _ = tree_flatten(inputs)
         model_output = executorch_module.run_method("forward", tuple(inputs_flattened))
         ref_output = m(*inputs)

exir/backend/test/test_backends_lifted.py

@@ -297,7 +297,6 @@ def forward(self, a, x, b):
 
         executorch_module = _load_for_executorch_from_buffer(buff)
 
-        # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`.
         inputs_flattened, _ = tree_flatten(model_inputs)
         model_output = executorch_module.run_method("forward", tuple(inputs_flattened))
         ref_output = add_mul_module(*model_inputs)
@@ -347,15 +346,17 @@ def forward(self, x):
 
         buff = exec_prog.buffer
 
+        executorch_module = _load_for_executorch_from_buffer(buff)
         # This line should raise an exception like
         # RuntimeError: failed with error 0x12
-        _load_for_executorch_from_buffer(buff)
+        inputs_flattened, _ = tree_flatten(model_inputs)
+        executorch_module.run_method("forward", tuple(inputs_flattened))
 
     @vary_segments
     def test_backend_with_compiler_out_of_range(self, extract_delegate_segments: bool):
         with self.assertRaisesRegex(
             RuntimeError,
-            "loading method forward failed with error 0x12",
+            "Failed to execute method forward, error: 0x12",
         ):
             self.run_model_in_unsupported_backend(
                 extract_delegate_segments=extract_delegate_segments
@@ -887,7 +888,6 @@ def forward(self, a, x, b):
         self.assertEqual(counter, 2)
 
         executorch_module = _load_for_executorch_from_buffer(executorch_prog.buffer)
-        # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`.
         inputs_flattened, _ = tree_flatten(inputs)
         model_output = executorch_module.run_method("forward", tuple(inputs_flattened))
         ref_output = m(*inputs)

exir/backend/test/test_compatibility.py

@@ -62,12 +62,13 @@ def forward(self, x):
             )
         )
 
+        executorch_module = _load_for_executorch_from_buffer(buff)
         # Throw runtime error with error code 0x30, meaning delegate is incompatible.
         with self.assertRaisesRegex(
             RuntimeError,
-            "loading method forward failed with error 0x30",
+            "Failed to execute method forward, error: 0x30",
         ):
-            executorch_module = _load_for_executorch_from_buffer(buff)
+            executorch_module.run_method("forward")
 
     def test_compatibility_in_runtime_edge_program_manager(self):
         class SinModule(torch.nn.Module):
@@ -108,9 +109,10 @@ def forward(self, x):
             )
         )
 
+        executorch_module = _load_for_executorch_from_buffer(buff)
         # Throw runtime error with error code 0x30, meaning delegate is incompatible.
         with self.assertRaisesRegex(
             RuntimeError,
-            "loading method forward failed with error 0x30",
+            "Failed to execute method forward, error: 0x30",
         ):
-            executorch_module = _load_for_executorch_from_buffer(buff)
+            executorch_module.run_method("forward")