Add input bucketization support for Gemini MM models.

PiperOrigin-RevId: 597078789 Change-Id: I8417622d5486418d333d1c015afe26960a1b2c49
google · Jan 10, 2024 · 26a056c · 26a056c
1 parent 51afaf8
commit 26a056c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/saxml/server/pax/BUILD b/saxml/server/pax/BUILD
@@ -28,7 +28,7 @@ pytype_strict_library(
     name = "branch_selection",
     srcs = ["branch_selection.py"],
     srcs_version = "PY3",
-    visibility = ["//saxml:internal"],
+    visibility = ["//visibility:public"],
     deps = [
         "//third_party/py/numpy",
         "//third_party/py/tensorflow:tensorflow_no_contrib",