test

PiperOrigin-RevId: 625513414
tensorflow · Apr 17, 2024 · 5eb46fe · 5eb46fe
1 parent 73851a0
commit 5eb46fe
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 11 deletions.
diff --git a/tensorflow_model_optimization/python/core/common/keras/compression/BUILD b/tensorflow_model_optimization/python/core/common/keras/compression/BUILD
@@ -1,4 +1,4 @@
-load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library", "py_strict_test", "pytype_strict_library")
+load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "pytype_strict_library", "pytype_strict_test")
 
 licenses(["notice"])
 
@@ -14,7 +14,7 @@ pytype_strict_library(
     ],
 )
 
-py_strict_library(
+pytype_strict_library(
     name = "schedules",
     srcs = ["schedules.py"],
     srcs_version = "PY3ONLY",
@@ -23,7 +23,7 @@ py_strict_library(
     ],
 )
 
-py_strict_test(
+pytype_strict_test(
     name = "schedules_test",
     srcs = [
         "schedules_test.py",

diff --git a/tensorflow_model_optimization/python/core/common/keras/compression/algorithms/BUILD b/tensorflow_model_optimization/python/core/common/keras/compression/algorithms/BUILD
@@ -1,4 +1,4 @@
-load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library", "py_strict_test", "pytype_strict_library")
+load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "pytype_strict_library", "pytype_strict_test")
 
 package(default_visibility = ["//visibility:private"])
 
@@ -16,7 +16,7 @@ pytype_strict_library(
     ],
 )
 
-py_strict_test(
+pytype_strict_test(
     name = "epr_test",
     timeout = "long",
     srcs = ["epr_test.py"],
@@ -46,7 +46,7 @@ pytype_strict_library(
     ],
 )
 
-py_strict_test(
+pytype_strict_test(
     name = "same_training_and_inference_test",
     timeout = "long",
     srcs = ["same_training_and_inference_test.py"],
@@ -72,7 +72,7 @@ pytype_strict_library(
     ],
 )
 
-py_strict_test(
+pytype_strict_test(
     name = "different_training_and_inference_test",
     timeout = "long",
     srcs = ["different_training_and_inference_test.py"],
@@ -98,7 +98,7 @@ pytype_strict_library(
     ],
 )
 
-py_strict_test(
+pytype_strict_test(
     name = "bias_only_test",
     timeout = "long",
     srcs = ["bias_only_test.py"],
@@ -113,7 +113,7 @@ py_strict_test(
     ],
 )
 
-py_strict_library(
+pytype_strict_library(
     name = "weight_clustering",
     srcs = ["weight_clustering.py"],
     srcs_version = "PY3",
@@ -127,7 +127,7 @@ py_strict_library(
     ],
 )
 
-py_strict_test(
+pytype_strict_test(
     name = "weight_clustering_test",
     srcs = ["weight_clustering_test.py"],
     python_version = "PY3",
@@ -152,7 +152,7 @@ pytype_strict_library(
     ],
 )
 
-py_strict_test(
+pytype_strict_test(
     name = "periodical_update_and_scheduling_test",
     timeout = "long",
     srcs = ["periodical_update_and_scheduling_test.py"],

diff --git a/tensorflow_model_optimization/tensorflow_model_optimization.bzl b/tensorflow_model_optimization/tensorflow_model_optimization.bzl
@@ -37,3 +37,7 @@ def pytype_library(name, **kwargs):
 # Placeholder to use until bazel supports py_strict_test.
 def py_strict_test(name, **kwargs):
     native.py_test(name = name, **kwargs)
+
+# Placeholder to use until bazel supports ptype_strict_test.
+def ptype_strict_test(name, **kwargs):
+    native.py_test(name = name, **kwargs)