testing

WORKSPACE

@@ -1,3 +1,18 @@
+# coding=utf-8
+# Copyright 2023 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """Workspace file for Jax Tpu Embedding."""
 
 load("//jax_tpu_embedding:build_defs.bzl", "jte_targets")  # @unused

jax_tpu_embedding/build_defs.bzl

@@ -1,6 +1,6 @@
 """Starlark macros for Jax Embedding API users."""
 
-load("//devtools/python/blaze:pytype.bzl", "pytype_binary")
+load("//jax_tpu_embedding:jte.bzl", "pytype_binary", "pytype_strict_binary")
 
 def _export_sources_impl(ctx):
     files = []

jax_tpu_embedding/jte.bzl

@@ -1,3 +1,18 @@
+# coding=utf-8
+# Copyright 2022 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """Implements custom rules for Jax Tpu Embedding."""
 
 # Placeholder to use until bazel supports pytype_*.

jax_tpu_embedding/pip_package/Dockerfile

@@ -74,7 +74,6 @@ RUN pip3 install pip six numpy wheel setuptools mock 'future>=0.17.1'
 RUN mv jax_tpu_embedding/jax_tpu_embedding/pip_package /jax_tpu_embedding/
 RUN cd /jax_tpu_embedding && bash pip_package/build.sh
 
-
 WORKDIR /
 
 CMD ["/bin/bash"]

jax_tpu_embedding/pip_package/build.sh

@@ -63,4 +63,4 @@ bazel test  --test_output=all --test_verbose_timeout_warnings   --  jax_tpu_embe
 
 chmod u+x pip_package/build_pip_pkg.sh
 ./pip_package/build_pip_pkg.sh "$DST_DIR" ${PYTHON_VERSION}
-pip3 freeze > "${DST_DIR}/dependencies.txt"
+pip3 freeze > "${DST_DIR}/dependencies.txt"