diff --git a/python/tflite_micro/BUILD b/python/tflite_micro/BUILD
index 282e9db7..4780b9a2 100644
--- a/python/tflite_micro/BUILD
+++ b/python/tflite_micro/BUILD
@@ -78,6 +78,7 @@ py_library(
     visibility = ["//visibility:public"],
     deps = [
         requirement("numpy"),
+        "//tensorflow/lite/micro/tools:generate_test_for_model",
         "//tensorflow/lite/tools:flatbuffer_utils",
     ],
 )
@@ -132,6 +133,7 @@ py_package(
     # in the tflm tree.
     packages = [
         "python.tflite_micro",
+        "tensorflow.lite.micro.tools.generate_test_for_model",
         "tensorflow.lite.python",
         "tensorflow.lite.tools.flatbuffer_utils",
     ],
diff --git a/tensorflow/lite/micro/tools/BUILD b/tensorflow/lite/micro/tools/BUILD
index a85a7bab..2d1e1874 100644
--- a/tensorflow/lite/micro/tools/BUILD
+++ b/tensorflow/lite/micro/tools/BUILD
@@ -32,6 +32,7 @@ py_library(
     name = "generate_test_for_model",
     srcs = ["generate_test_for_model.py"],
     srcs_version = "PY3",
+    visibility = ["//:__subpackages__"],
     deps = [
         "//tensorflow/lite/python:schema_py",
     ],