build(bazel): inject tflite_micro shim via tflm_py_* wrappers

BUILD

@@ -21,6 +21,11 @@ config_setting(
     },
 )
 
+# Synthesizes the "tflite_micro" top-level package namespace at import time, so
+# that `from tflite_micro...` imports resolve (see tflite_micro.py). Necessary
+# because Bzlmod fixes the main repo's runfiles root to "_main" rather than the
+# module name. Targets should generally not depend on this directly: use the
+# tflm_py_* wrappers in //python:py_rules.bzl, which add it automatically.
 py_library(
     name = "tflite_micro_shim",
     srcs = ["tflite_micro.py"],

python/py_rules.bzl

@@ -0,0 +1,37 @@
+"""Bazel macros wrapping @rules_python's py_* rules for use in this repository.
+
+The tflm_py_library, tflm_py_test, and tflm_py_binary wrappers inject a
+dependency on //:tflite_micro_shim, which synthesizes the "tflite_micro"
+top-level package namespace at import time (see //:tflite_micro.py).
+
+The shim is necessary because, under Bzlmod, the main repository's runfiles root
+is the fixed name "_main" rather than the module name, so the "tflite_micro"
+import prefix used throughout this repository's Python sources (e.g. `from
+tflite_micro.tensorflow.lite... import x`) no longer resolves on its own as it
+did under the legacy WORKSPACE, where the runfiles root took the workspace name.
+
+Any py_* target whose code imports "tflite_micro.*" must depend on the shim.
+These wrappers add that dependency automatically so it need not be repeated in
+every target. Prefer them over the native py_library/py_test/py_binary for any
+target under the tflite_micro namespace.
+"""
+
+load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
+
+def tflm_py_library(deps = [], **kwargs):
+    py_library(
+        deps = deps + ["//:tflite_micro_shim"],
+        **kwargs
+    )
+
+def tflm_py_binary(deps = [], **kwargs):
+    py_binary(
+        deps = deps + ["//:tflite_micro_shim"],
+        **kwargs
+    )
+
+def tflm_py_test(deps = [], **kwargs):
+    py_test(
+        deps = deps + ["//:tflite_micro_shim"],
+        **kwargs
+    )

python/tflite_micro/BUILD

@@ -5,6 +5,7 @@ load("@rules_python//python:packaging.bzl", "py_package", "py_wheel")
 load("@rules_shell//shell:sh_test.bzl", "sh_test")  # @unused
 load("@tflm_pip_deps//:requirements.bzl", "requirement")
 load("//python:py_namespace.bzl", "py_namespace")
+load("//python:py_rules.bzl", "tflm_py_library")
 load(
     "//tensorflow:extra_rules.bzl",
     "tflm_python_op_resolver_friends",
@@ -71,7 +72,7 @@ pybind_extension(
     ],
 )
 
-py_library(
+tflm_py_library(
     name = "runtime",
     srcs = [
         "runtime.py",
@@ -82,7 +83,6 @@ py_library(
     target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
     visibility = ["//visibility:public"],
     deps = [
-        "//:tflite_micro_shim",
         requirement("numpy"),
         "//tensorflow/lite/micro/tools:generate_test_for_model",
         "//tensorflow/lite/python:schema_py",

python/tflite_micro/signal/utils/BUILD

@@ -3,6 +3,7 @@ load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
 # Signal python utilities.
 load("@rules_python//python:defs.bzl", "py_library", "py_test")
 load("@tflm_pip_deps//:requirements.bzl", "requirement")
+load("//python:py_rules.bzl", "tflm_py_library", "tflm_py_test")
 load("//tensorflow/lite/micro:build_def.bzl", "INCOMPATIBLE_WITH_WINDOWS")
 
 package(
@@ -12,7 +13,7 @@ package(
     licenses = ["notice"],
 )
 
-py_test(
+tflm_py_test(
     name = "freq_to_mel_test",
     srcs = ["freq_to_mel_test.py"],
     data = [
@@ -24,20 +25,14 @@ py_test(
         "noubsan",
     ],
     target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
-    deps = [
-        "//:tflite_micro_shim",
-    ],
 )
 
-py_library(
+tflm_py_library(
     name = "freq_to_mel",
     data = [
         ":freq_to_mel_wrapper.so",
     ],
     target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
-    deps = [
-        "//:tflite_micro_shim",
-    ],
 )
 
 pybind_extension(
@@ -67,18 +62,15 @@ pybind_extension(
     target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
 )
 
-py_library(
+tflm_py_library(
     name = "wide_dynamic_func_lut",
     data = [
         ":wide_dynamic_func_lut_wrapper.so",
     ],
     target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
-    deps = [
-        "//:tflite_micro_shim",
-    ],
 )
 
-py_test(
+tflm_py_test(
     name = "wide_dynamic_func_lut_test",
     srcs = ["wide_dynamic_func_lut_test.py"],
     data = [
@@ -91,7 +83,4 @@ py_test(
     ],
     target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
     visibility = ["//visibility:public"],
-    deps = [
-        "//:tflite_micro_shim",
-    ],
 )

tensorflow/lite/micro/compression/BUILD

@@ -1,6 +1,12 @@
 load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
 load("@rules_shell//shell:sh_test.bzl", "sh_test")
 load("@tflm_pip_deps//:requirements.bzl", "requirement")
+load(
+    "//python:py_rules.bzl",
+    "tflm_py_binary",
+    "tflm_py_library",
+    "tflm_py_test",
+)
 load(
     "//tensorflow/lite/micro:build_def.bzl",
     "INCOMPATIBLE_WITH_WINDOWS",
@@ -99,15 +105,14 @@ flatbuffer_py_library(
     srcs = ["metadata.fbs"],
 )
 
-py_test(
+tflm_py_test(
     name = "metadata_test_py",
     size = "small",
     srcs = ["metadata_test.py"],
     main = "metadata_test.py",
     target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
     deps = [
         "metadata_py",
-        "//:tflite_micro_shim",
         requirement("flatbuffers"),
     ],
 )
@@ -164,11 +169,10 @@ py_test(
     ],
 )
 
-py_library(
+tflm_py_library(
     name = "model_facade",
     srcs = ["model_facade.py"],
     deps = [
-        "//:tflite_micro_shim",
         "//tensorflow/lite/python:schema_py",
         requirement("flatbuffers"),
     ],
@@ -185,11 +189,10 @@ py_test(
     ],
 )
 
-py_library(
+tflm_py_library(
     name = "spec",
     srcs = ["spec.py"],
     deps = [
-        "//:tflite_micro_shim",
         requirement("pyyaml"),
     ],
 )
@@ -223,11 +226,10 @@ py_test(
     ],
 )
 
-py_library(
+tflm_py_library(
     name = "test_models",
     srcs = ["test_models.py"],
     deps = [
-        "//:tflite_micro_shim",
         "//tensorflow/lite/python:schema_py",
         requirement("flatbuffers"),
         requirement("numpy"),
@@ -245,15 +247,14 @@ py_test(
     ],
 )
 
-py_binary(
+tflm_py_binary(
     name = "view",
     srcs = [
         "view.py",
     ],
     target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
     deps = [
         ":metadata_py",
-        "//:tflite_micro_shim",
         "//tensorflow/lite/python:schema_py",
         requirement("absl_py"),
         requirement("bitarray"),

tensorflow/lite/micro/examples/person_detection/utils/BUILD

@@ -1,5 +1,6 @@
 load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
 load("@tflm_pip_deps//:requirements.bzl", "requirement")
+load("//python:py_rules.bzl", "tflm_py_test")
 
 package(
     features = ["-layering_check"],
@@ -22,7 +23,7 @@ py_library(
     ],
 )
 
-py_test(
+tflm_py_test(
     name = "raw_to_bitmap_test",
     srcs = ["raw_to_bitmap_test.py"],
     data = glob(
@@ -36,7 +37,6 @@ py_test(
     ],
     deps = [
         ":raw_to_bitmap_lib",
-        "//:tflite_micro_shim",
         requirement("numpy"),
     ],
 )

tensorflow/lite/micro/kernels/testdata/BUILD

@@ -1,5 +1,6 @@
 load("@rules_python//python:defs.bzl", "py_binary", "py_test")
 load("@tflm_pip_deps//:requirements.bzl", "requirement")
+load("//python:py_rules.bzl", "tflm_py_binary")
 load(
     "//tensorflow/lite/micro:build_def.bzl",
     "tflm_cc_library",
@@ -39,14 +40,13 @@ tflm_cc_library(
 ####################################
 # Python
 ####################################
-py_binary(
+tflm_py_binary(
     name = "lstm_test_data_generator",
     srcs = [
         "lstm_test_data_generator.py",
         "lstm_test_data_utils.py",
     ],
     deps = [
-        "//:tflite_micro_shim",
         requirement("absl_py"),
         requirement("numpy"),
         requirement("tensorflow"),

tensorflow/lite/micro/python/tflite_size/tests/BUILD

@@ -1,5 +1,5 @@
-load("@rules_python//python:defs.bzl", "py_test")
 load("@tflm_pip_deps//:requirements.bzl", "requirement")
+load("//python:py_rules.bzl", "tflm_py_test")
 load("//tensorflow/lite/micro:build_def.bzl", "INCOMPATIBLE_WITH_WINDOWS")
 
 licenses(["notice"])
@@ -13,7 +13,7 @@ filegroup(
     ],
 )
 
-py_test(
+tflm_py_test(
     name = "flatbuffer_size_test",
     srcs = ["flatbuffer_size_test.py"],
     data = [
@@ -27,7 +27,6 @@ py_test(
     ],
     target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
     deps = [
-        "//:tflite_micro_shim",
         "//tensorflow/lite/micro/python/tflite_size/src:flatbuffer_size",
     ],
 )

tensorflow/lite/micro/tools/BUILD

@@ -2,11 +2,14 @@ load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
 load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
 load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
 load("@tflm_pip_deps//:requirements.bzl", "requirement")
-load("//tensorflow:extra_rules.bzl", "tflm_application_friends")
 load(
-    "//tensorflow/lite/micro:build_def.bzl",
-    "INCOMPATIBLE_WITH_WINDOWS",
+    "//python:py_rules.bzl",
+    "tflm_py_binary",
+    "tflm_py_library",
+    "tflm_py_test",
 )
+load("//tensorflow:extra_rules.bzl", "tflm_application_friends")
+load("//tensorflow/lite/micro:build_def.bzl", "INCOMPATIBLE_WITH_WINDOWS")
 load("//third_party/flatbuffers:build_defs.bzl", "flatbuffer_cc_library", "flatbuffer_py_library")
 
 package(
@@ -33,12 +36,11 @@ py_library(
     ],
 )
 
-py_library(
+tflm_py_library(
     name = "generate_test_for_model",
     srcs = ["generate_test_for_model.py"],
     visibility = ["//:__subpackages__"],
     deps = [
-        "//:tflite_micro_shim",
         "//tensorflow/lite/python:schema_py",
         requirement("ai-edge-litert"),
     ],
@@ -53,14 +55,13 @@ py_binary(
     ],
 )
 
-py_binary(
+tflm_py_binary(
     name = "requantize_flatbuffer",
     srcs = [
         "requantize_flatbuffer.py",
         "requantize_flatbuffer_utils.py",
     ],
     deps = [
-        "//:tflite_micro_shim",
         "//tensorflow/lite/python:schema_py",
         "//tensorflow/lite/tools:flatbuffer_utils",
         requirement("absl_py"),
@@ -111,7 +112,7 @@ py_binary(
     ],
 )
 
-py_library(
+tflm_py_library(
     name = "model_transforms_utils",
     srcs = ["model_transforms_utils.py"],
     data = [
@@ -123,7 +124,6 @@ py_library(
         ":tflm_tools",
     ],
     deps = [
-        "//:tflite_micro_shim",
         "//tensorflow/lite/python:schema_py",
         "//tensorflow/lite/python:schema_util",
     ],
@@ -183,7 +183,7 @@ py_binary(
     ],
 )
 
-py_test(
+tflm_py_test(
     name = "tflm_model_transforms_test",
     srcs = ["tflm_model_transforms_test.py"],
     data = [
@@ -198,7 +198,6 @@ py_test(
     target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
     deps = [
         ":tflm_model_transforms_lib",
-        "//:tflite_micro_shim",
         requirement("absl_py"),
         requirement("tensorflow"),
         "//tensorflow/lite/micro/examples/recipes:resource_variables_lib",

tensorflow/lite/tools/BUILD

@@ -1,12 +1,12 @@
 load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
 load("@tflm_pip_deps//:requirements.bzl", "requirement")
+load("//python:py_rules.bzl", "tflm_py_library")
 
-py_library(
+tflm_py_library(
     name = "flatbuffer_utils",
     srcs = ["flatbuffer_utils.py"],
     visibility = ["//:__subpackages__"],
     deps = [
-        "//:tflite_micro_shim",
         requirement("flatbuffers"),
         requirement("tensorflow"),
         "//tensorflow/lite/python:schema_py",