Check in generated pyi files for some py_extension targets.

PiperOrigin-RevId: 554878889

tensorflow/lite/python/analyzer_wrapper/BUILD

@@ -11,6 +11,10 @@ pybind_extension(
     srcs = [
         "analyzer_wrapper.cc",
     ],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_pywrap_analyzer_wrapper.pyi",
+    ],
     deps = [
         ":model_analyzer",
         "@pybind11",

tensorflow/lite/python/analyzer_wrapper/_pywrap_analyzer_wrapper.pyi

@@ -0,0 +1,16 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+def ModelAnalyzer(arg0: str, arg1: bool, arg2: bool) -> str: ...

tensorflow/lite/python/metrics/BUILD

@@ -37,6 +37,10 @@ pybind_extension(
     srcs = ["wrapper/metrics_wrapper_pybind11.cc"],
     hdrs = ["wrapper/metrics_wrapper.h"],
     compatible_with = get_compatible_with_portable(),
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_pywrap_tensorflow_lite_metrics_wrapper.pyi",
+    ],
     visibility = ["//visibility:private"],
     deps = [
         ":metrics_wrapper_lib",

tensorflow/lite/python/metrics/_pywrap_tensorflow_lite_metrics_wrapper.pyi

@@ -0,0 +1,18 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+class MetricsWrapper:
+    def __init__(self, arg0: str) -> None: ...
+    def ExportMetrics(self) -> object: ...

tensorflow/lite/python/testdata/BUILD

@@ -175,7 +175,11 @@ pybind_extension(
     ],
     hdrs = ["test_registerer.h"],
     additional_exported_symbols = ["TF_TestRegisterer"],
+    enable_stub_generation = True,
     link_in_framework = True,
+    pytype_srcs = [
+        "_pywrap_test_registerer.pyi",
+    ],
     deps = [
         ":test_registerer",
         "//tensorflow/lite:framework",

tensorflow/lite/python/testdata/_pywrap_test_registerer.pyi

@@ -0,0 +1,17 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+def TF_TestRegisterer(arg0: int) -> None: ...
+def get_num_test_registerer_calls() -> int: ...

tensorflow/lite/tools/optimize/python/BUILD

@@ -57,6 +57,10 @@ py_strict_library(
 pybind_extension(
     name = "_pywrap_modify_model_interface",
     srcs = ["modify_model_interface.cc"],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_pywrap_modify_model_interface.pyi",
+    ],
     deps = [
         "//tensorflow/lite/schema:schema_fbs",
         "//tensorflow/lite/tools/optimize:modify_model_interface",

tensorflow/lite/tools/optimize/python/_pywrap_modify_model_interface.pyi

@@ -0,0 +1,16 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+def modify_model_interface(arg0: str, arg1: str, arg2: int, arg3: int) -> int: ...