add typing to pybindings

backends/qnnpack/test/TARGETS

@@ -16,7 +16,6 @@ python_unittest(
         "//executorch/backends/qnnpack/partition:qnnpack_partitioner",
         "//executorch/exir:lib",
         "//executorch/exir/backend:backend_api",
-        "//executorch/exir/serialize:lib",
         "//executorch/extension/pybindings:portable",  # @manual
         "//executorch/extension/pytree:pylib",
     ],

backends/qnnpack/test/test_qnnpack.py

@@ -19,7 +19,6 @@
 
 from executorch.exir.backend.backend_api import to_backend, validation_disabled
 
-# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`.
 from executorch.extension.pybindings.portable import (  # @manual
     _load_for_executorch_from_buffer,
 )

backends/vulkan/test/test_vulkan_delegate.py

@@ -19,7 +19,7 @@
 
 ctypes.CDLL("libvulkan.so.1")
 
-# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`.
+
 from executorch.extension.pybindings.portable import (  # @manual
     _load_for_executorch_from_buffer,
 )
@@ -85,7 +85,6 @@ def forward(self, *args):
         )
 
         # Test the model with executor
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         executorch_module = _load_for_executorch_from_buffer(executorch_program.buffer)
         # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`.
         inputs_flattened, _ = tree_flatten(sample_inputs)

backends/xnnpack/test/TARGETS

@@ -28,7 +28,6 @@ python_unittest(
         "//executorch/exir:tracer",
         "//executorch/exir/backend:backend_api",
         "//executorch/exir/passes:spec_prop_pass",
-        "//executorch/exir/serialize:lib",
         "//executorch/extension/pybindings:portable",  # @manual
         "//executorch/extension/pytree:pylib",
     ],
@@ -59,7 +58,6 @@ python_unittest(
         "//executorch/exir/backend:backend_api",
         "//executorch/exir/dialects:lib",
         "//executorch/exir/passes:spec_prop_pass",
-        "//executorch/exir/serialize:lib",
         "//executorch/extension/pybindings:portable",  # @manual
         "//executorch/extension/pytree:pylib",
     ],
@@ -89,7 +87,6 @@ python_unittest(
         "//executorch/exir:tracer",
         "//executorch/exir/backend:backend_api",
         "//executorch/exir/passes:spec_prop_pass",
-        "//executorch/exir/serialize:lib",
         "//executorch/extension/pybindings:portable",  # @manual
         "//executorch/extension/pytree:pylib",
         "//pytorch/vision:torchvision",
@@ -133,6 +130,7 @@ python_unittest(
         "//caffe2:torch",
         "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
         "//executorch/backends/xnnpack/test/tester:tester",
+        "//executorch/backends/xnnpack/utils:xnnpack_utils",
         "//pytorch/vision:torchvision",
     ],
 )

backends/xnnpack/test/test_xnnpack_utils.py

@@ -37,7 +37,6 @@
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
 from executorch.exir.tracer import _default_decomposition_table
 
-# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`.
 from executorch.extension.pybindings.portable import (  # @manual
     _load_for_executorch_from_buffer,
 )
@@ -230,7 +229,6 @@ def forward(self, *args):
             )
 
         # Test the model with executor
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         executorch_module = _load_for_executorch_from_buffer(executorch_program.buffer)
         # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`.
         inputs_flattened, _ = tree_flatten(sample_inputs)
@@ -439,7 +437,6 @@ def forward(self, x):
                 output_path=filename,
             )
 
-        # pyre-ignore
         executorch_module = _load_for_executorch_from_buffer(executorch_program.buffer)
         # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`.
         inputs_flattened, _ = tree_flatten(example_inputs)

backends/xnnpack/test/tester/tester.py

@@ -30,10 +30,7 @@
 from executorch.exir.backend.partitioner import Partitioner
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
 
-# pyre-ignore[21]: Could not find module `executorch.pybindings.portable`.
-from executorch.extension.pybindings.portable import (  # @manual
-    _load_for_executorch_from_buffer,
-)
+from executorch.extension.pybindings.portable import _load_for_executorch_from_buffer
 from torch.ao.quantization.backend_config import BackendConfig
 from torch.ao.quantization.backend_config.executorch import (
     get_executorch_backend_config,

examples/export/test/test_export.py

@@ -13,7 +13,6 @@
 from executorch.examples.export.utils import export_to_edge
 from executorch.examples.models import MODEL_NAME_TO_MODEL
 
-# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`.
 from executorch.extension.pybindings.portable import (  # @manual
     _load_for_executorch_from_buffer,
 )
@@ -37,7 +36,7 @@ def _assert_eager_lowered_same_result(
         edge_model = export_to_edge(eager_model, example_inputs)
 
         executorch_prog = edge_model.to_executorch()
-        # pyre-ignore
+
         pte_model = _load_for_executorch_from_buffer(executorch_prog.buffer)
 
         with torch.no_grad():

exir/backend/test/demos/rpc/test_rpc.py

@@ -17,7 +17,6 @@
 )
 from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo
 
-# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`.
 from executorch.extension.pybindings.portable import (  # @manual
     _load_for_executorch_from_buffer,
 )

exir/backend/test/demos/test_delegate_aten_mode.py

@@ -15,7 +15,6 @@
     BackendWithCompilerDemo,
 )
 
-# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`.
 from executorch.extension.pybindings.aten_mode_lib import (  # @manual
     _load_for_executorch_from_buffer,
 )

exir/backend/test/demos/test_xnnpack_qnnpack.py

@@ -22,7 +22,6 @@
 from executorch.exir.backend.backend_api import to_backend, validation_disabled
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
 
-# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`.
 from executorch.extension.pybindings.portable import (  # @manual
     _load_for_executorch_from_buffer,
 )

exir/backend/test/test_backends.py

@@ -46,7 +46,6 @@
     Program,
 )
 
-# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`.
 from executorch.extension.pybindings.portable import (  # @manual
     _load_for_executorch_from_buffer,
 )
@@ -224,7 +223,6 @@ def forward(self, x):
         )
         buff = exec_prog.buffer
 
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         executorch_module = _load_for_executorch_from_buffer(buff)
         model_inputs = torch.ones(1)
         model_outputs = executorch_module.forward([model_inputs])
@@ -281,7 +279,6 @@ def forward(self, a, x, b):
         )
         buff = exec_prog.buffer
 
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         executorch_module = _load_for_executorch_from_buffer(buff)
 
         # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`.
@@ -338,7 +335,6 @@ def forward(self, x):
 
         # This line should raise an exception like
         # RuntimeError: failed with error 0x12
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         _load_for_executorch_from_buffer(buff)
 
     @vary_segments
@@ -434,7 +430,6 @@ def forward(self, x):
             )
         )
 
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         executorch_module = _load_for_executorch_from_buffer(buff)
         model_inputs = torch.ones(1)
 
@@ -561,7 +556,6 @@ def forward(self, x):
         )
         flatbuffer = exec_prog.buffer
 
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         executorch_module = _load_for_executorch_from_buffer(flatbuffer)
         model_outputs = executorch_module.forward([*model_inputs])
 
@@ -858,7 +852,6 @@ def forward(self, a, x, b):
         # There should be 2 delegated modules
         self.assertEqual(counter, 2)
 
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         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)

exir/backend/test/test_backends_lifted.py

@@ -49,7 +49,6 @@
     Program,
 )
 
-# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`.
 from executorch.extension.pybindings.portable import (  # @manual
     _load_for_executorch_from_buffer,
 )
@@ -231,7 +230,6 @@ def forward(self, x):
         )
         buff = exec_prog.buffer
 
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         executorch_module = _load_for_executorch_from_buffer(buff)
         model_inputs = torch.ones(1)
         model_outputs = executorch_module.forward([model_inputs])
@@ -290,7 +288,6 @@ def forward(self, a, x, b):
         )
         buff = exec_prog.buffer
 
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         executorch_module = _load_for_executorch_from_buffer(buff)
 
         # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`.
@@ -347,7 +344,6 @@ def forward(self, x):
 
         # This line should raise an exception like
         # RuntimeError: failed with error 0x12
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         _load_for_executorch_from_buffer(buff)
 
     @vary_segments
@@ -443,7 +439,6 @@ def forward(self, x):
             )
         )
 
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         executorch_module = _load_for_executorch_from_buffer(buff)
         model_inputs = torch.ones(1)
 
@@ -570,7 +565,6 @@ def forward(self, x):
         )
         flatbuffer = exec_prog.buffer
 
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         executorch_module = _load_for_executorch_from_buffer(flatbuffer)
         model_outputs = executorch_module.forward([*model_inputs])
 
@@ -869,7 +863,6 @@ def forward(self, a, x, b):
         # There should be 2 delegated modules
         self.assertEqual(counter, 2)
 
-        # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
         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)

exir/emit/test/test_emit.py

@@ -38,10 +38,7 @@
 from executorch.exir.tests.common import register_additional_test_aten_ops
 from executorch.exir.tests.models import MLP, Mul
 
-# pyre-ignore
-from executorch.extension.pybindings.portable import (  # @manual
-    _load_for_executorch_from_buffer,
-)
+from executorch.extension.pybindings.portable import _load_for_executorch_from_buffer
 from functorch.experimental import control_flow
 
 

extension/pybindings/TARGETS

@@ -3,10 +3,25 @@
 # targets.bzl. This file can contain fbcode-only targets.
 
 load("@fbcode//executorch/extension/pybindings:targets.bzl", "ATEN_MODULE_DEPS", "MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB", "MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB", "PORTABLE_MODULE_DEPS", "define_common_targets", "executorch_pybindings")
+load("@fbcode_macros//build_defs:native_rules.bzl", "buck_genrule")
 load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
 
 define_common_targets()
 
+# In order to have pyre recognize the pybindings module, the name of the .pyi must exactly match the
+# name of the lib. To avoid copy pasting the pyi file in tree a whole bunch of times we use genrules
+# to do it for us
+buck_genrule(
+    name = "pybindings_types_gen",
+    srcs = [":pybinding_types"],
+    outs = {
+        "aten_mode_lib.pyi": ["aten_mode_lib.pyi"],
+        "portable.pyi": ["portable.pyi"],
+    },
+    cmd = "cp $(location :pybinding_types)/* $OUT/portable.pyi && cp $(location :pybinding_types)/* $OUT/aten_mode_lib.pyi",
+    visibility = ["//executorch/extension/pybindings/..."],
+)
+
 executorch_pybindings(
     srcs = [
         "module.cpp",
@@ -22,6 +37,7 @@ executorch_pybindings(
     ],
     cppdeps = PORTABLE_MODULE_DEPS + MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB,
     python_module_name = "portable",
+    types = ["//executorch/extension/pybindings:pybindings_types_gen[portable.pyi]"],
     visibility = ["PUBLIC"],
 )
 
@@ -31,6 +47,7 @@ executorch_pybindings(
     ],
     cppdeps = ATEN_MODULE_DEPS + MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB,
     python_module_name = "aten_mode_lib",
+    types = ["//executorch/extension/pybindings:pybindings_types_gen[aten_mode_lib.pyi]"],
     visibility = ["PUBLIC"],
 )
 

extension/pybindings/module.cpp

@@ -403,15 +403,17 @@ struct PyModule final {
     return std::make_unique<PyModule>(m.get_program_ptr(), m.get_program_len());
   }
 
-  py::list run_method(const std::string& name, const py::sequence& pyinputs) {
-    std::vector<EValue> inputs;
-    const auto inputs_size = py::len(pyinputs);
-    inputs.reserve(inputs_size);
+  py::list run_method(
+      const std::string& method_name,
+      const py::sequence& inputs) {
+    std::vector<EValue> cpp_inputs;
+    const auto inputs_size = py::len(inputs);
+    cpp_inputs.reserve(inputs_size);
     for (size_t i = 0; i < inputs_size; ++i) {
-      inputs.emplace_back(pyToEValue(pyinputs[i], keep_alive_));
+      cpp_inputs.emplace_back(pyToEValue(inputs[i], keep_alive_));
     }
 
-    auto outputs = module_->run_method(name, inputs);
+    auto outputs = module_->run_method(method_name, cpp_inputs);
 
     const auto outputs_size = outputs.size();
     py::list list(outputs_size);
@@ -421,8 +423,8 @@ struct PyModule final {
     return list;
   }
 
-  py::list forward(const py::sequence& pyinputs) {
-    return run_method("forward", pyinputs);
+  py::list forward(const py::sequence& inputs) {
+    return run_method("forward", inputs);
   }
 
  private:
@@ -461,7 +463,7 @@ void init_module_functions(py::module_& m) {
   m.def("_create_profile_block", &create_profile_block);
   m.def("_reset_profile_results", []() { EXECUTORCH_RESET_PROFILE_RESULTS(); });
 
-  py::class_<PyModule>(m, "Module")
+  py::class_<PyModule>(m, "ExecutorchModule")
       .def("run_method", &PyModule::run_method)
       .def("forward", &PyModule::forward);
 

extension/pybindings/pybindings.pyi

@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+from typing import Any, Dict, List, Sequence, Tuple
+
+class ExecutorchModule:
+    def run_method(self, method_name: str, inputs: Sequence[Any]) -> List[Any]: ...
+    def forward(self, inputs: Sequence[Any]) -> List[Any]: ...
+
+def _load_for_executorch(path: str) -> ExecutorchModule: ...
+def _load_for_executorch_from_buffer(buffer: bytes) -> ExecutorchModule: ...
+def _create_profile_block(name: str) -> None: ...
+def _dump_profile_results() -> bytes: ...
+def _reset_profile_results() -> None: ...