From 002834de05685cced6372035a19946756dfc80b5 Mon Sep 17 00:00:00 2001 From: Ryan Kuester Date: Sun, 24 May 2026 22:40:44 -0500 Subject: [PATCH] build(bazel): inject tflite_micro shim via tflm_py_* wrappers Python targets in this repository import one another under the "tflite_micro" package namespace, which //:tflite_micro_shim synthesizes at import time. The shim is required under Bzlmod, where the main repository's runfiles root is the fixed name "_main" rather than the module name, so the "tflite_micro" prefix no longer resolves on its own. Every such target therefore had to list //:tflite_micro_shim in its deps, which was repetitive and easy to forget. Add tflm_py_library, tflm_py_test, and tflm_py_binary wrappers in a new //python:py_rules.bzl that inject the shim dependency automatically, following the naming convention of the existing tflm_cc_* wrappers, and document the shim's rationale there. Convert every target that previously listed the shim to the corresponding wrapper and drop the explicit dependency. The dependency graph is unchanged; only the means by which the shim is attached differs. --- BUILD | 5 +++ python/py_rules.bzl | 37 +++++++++++++++++++ python/tflite_micro/BUILD | 4 +- python/tflite_micro/signal/utils/BUILD | 21 +++-------- tensorflow/lite/micro/compression/BUILD | 21 ++++++----- .../examples/person_detection/utils/BUILD | 4 +- tensorflow/lite/micro/kernels/testdata/BUILD | 4 +- .../lite/micro/python/tflite_size/tests/BUILD | 5 +-- tensorflow/lite/micro/tools/BUILD | 21 +++++------ tensorflow/lite/tools/BUILD | 4 +- tools/BUILD | 4 +- 11 files changed, 80 insertions(+), 50 deletions(-) create mode 100644 python/py_rules.bzl diff --git a/BUILD b/BUILD index 7e4516056ff..c910a99dc4a 100644 --- a/BUILD +++ b/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"], diff --git a/python/py_rules.bzl b/python/py_rules.bzl new file mode 100644 index 00000000000..d67ac44af7d --- /dev/null +++ b/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 + ) diff --git a/python/tflite_micro/BUILD b/python/tflite_micro/BUILD index 0fdf5b53428..b358fd12adc 100644 --- a/python/tflite_micro/BUILD +++ b/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", diff --git a/python/tflite_micro/signal/utils/BUILD b/python/tflite_micro/signal/utils/BUILD index 05e204f75c3..fdd00b54f3f 100644 --- a/python/tflite_micro/signal/utils/BUILD +++ b/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", - ], ) diff --git a/tensorflow/lite/micro/compression/BUILD b/tensorflow/lite/micro/compression/BUILD index e8ecb2ba75d..408fc33912e 100644 --- a/tensorflow/lite/micro/compression/BUILD +++ b/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,7 +105,7 @@ flatbuffer_py_library( srcs = ["metadata.fbs"], ) -py_test( +tflm_py_test( name = "metadata_test_py", size = "small", srcs = ["metadata_test.py"], @@ -107,7 +113,6 @@ py_test( 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,7 +247,7 @@ py_test( ], ) -py_binary( +tflm_py_binary( name = "view", srcs = [ "view.py", @@ -253,7 +255,6 @@ py_binary( target_compatible_with = INCOMPATIBLE_WITH_WINDOWS, deps = [ ":metadata_py", - "//:tflite_micro_shim", "//tensorflow/lite/python:schema_py", requirement("absl_py"), requirement("bitarray"), diff --git a/tensorflow/lite/micro/examples/person_detection/utils/BUILD b/tensorflow/lite/micro/examples/person_detection/utils/BUILD index 3cbefd4da54..b059fef6865 100644 --- a/tensorflow/lite/micro/examples/person_detection/utils/BUILD +++ b/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"), ], ) diff --git a/tensorflow/lite/micro/kernels/testdata/BUILD b/tensorflow/lite/micro/kernels/testdata/BUILD index feb6c52e5d5..f268f15b6f6 100644 --- a/tensorflow/lite/micro/kernels/testdata/BUILD +++ b/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"), diff --git a/tensorflow/lite/micro/python/tflite_size/tests/BUILD b/tensorflow/lite/micro/python/tflite_size/tests/BUILD index b53854c468b..7b13c8b6c9c 100644 --- a/tensorflow/lite/micro/python/tflite_size/tests/BUILD +++ b/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", ], ) diff --git a/tensorflow/lite/micro/tools/BUILD b/tensorflow/lite/micro/tools/BUILD index caa6bd83ddd..bfa6bc4302e 100644 --- a/tensorflow/lite/micro/tools/BUILD +++ b/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", diff --git a/tensorflow/lite/tools/BUILD b/tensorflow/lite/tools/BUILD index 6316a167e3b..0b6ecaa62a3 100644 --- a/tensorflow/lite/tools/BUILD +++ b/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", diff --git a/tools/BUILD b/tools/BUILD index 204842e2bb8..7e9161d1b14 100644 --- a/tools/BUILD +++ b/tools/BUILD @@ -1,4 +1,5 @@ load("@rules_python//python:defs.bzl", "py_binary", "py_test") +load("//python:py_rules.bzl", "tflm_py_test") package( default_visibility = ["//visibility:public"], @@ -9,11 +10,10 @@ py_binary( srcs = ["expand_stamp_vars.py"], ) -py_test( +tflm_py_test( name = "expand_stamp_vars_test", srcs = ["expand_stamp_vars_test.py"], deps = [ ":expand_stamp_vars", - "//:tflite_micro_shim", ], )