Added Python toolchain setup in WORKSPACE.

TF moved Python toolchain setup from bzl scripts to the main WORKSPACE.
However, those bzl scripts (which we use for C++ toolchain setup)
expect the Python toolchain setup as done in TF WORKSPACE.

PiperOrigin-RevId: 572960752
Change-Id: I4ee41e6c24a6be584efbeacf4f30ee4db3075b7d

Author: ssjhv

README.md

@@ -269,7 +269,10 @@ sudo docker run -i --rm -v /tmp/tensorflow_compression:/tmp/tensorflow_compressi
      cd /tensorflow_compression &&
      python -m pip install -U pip setuptools wheel &&
      python -m pip install -r requirements.txt &&
-     bazel run -c opt --copt=-mavx --crosstool_top=@ubuntu20.04-gcc9_manylinux2014-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain :build_pip_pkg -- . /tmp/tensorflow_compression <custom-version>"
+     bazel build -c opt --copt=-mavx --crosstool_top=@ubuntu20.04-gcc9_manylinux2014-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain :build_pip_pkg &&
+     pushd bazel-bin/build_pip_pkg.runfiles/tensorflow_compression &&
+     python build_pip_pkg.py . /tmp/tensorflow_compression <custom-version> &&
+     popd"
 ```
 
 For Darwin, the Docker image and specifying the toolchain is not necessary. We
@@ -281,7 +284,10 @@ git clone https://github.com/tensorflow/compression.git /tensorflow_compression
 cd /tensorflow_compression
 python -m pip install -U pip setuptools wheel
 python -m pip install -r requirements.txt
-bazel run -c opt --copt=-mavx --macos_minimum_os=10.14 :build_pip_pkg -- . /tmp/tensorflow_compression <custom-version>"
+bazel build -c opt --copt=-mavx --macos_minimum_os=10.14 :build_pip_pkg
+pushd bazel-bin/build_pip_pkg.runfiles/tensorflow_compression
+python build_pip_pkg.py . /tmp/tensorflow_compression <custom-version>"
+popd
 ```
 
 In both cases, the wheel file is created inside `/tmp/tensorflow_compression`.

WORKSPACE

@@ -1,9 +1,28 @@
 workspace(name = "tensorflow_compression")
 
-load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load("//:workspace.bzl", "tensorflow_compression_workspace")
 tensorflow_compression_workspace()
 
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+# `bazel_skylib` and `rules_python` versions should match the ones used in
+# `org_tensorflow`.
+http_archive(
+    name = "bazel_skylib",
+    sha256 = "74d544d96f4a5bb630d465ca8bbcfe231e3594e5aae57e1edbf17a6eb3ca2506",
+    urls = [
+        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
+        "https://github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
+    ],
+)
+
+http_archive(
+    name = "rules_python",
+    sha256 = "9d04041ac92a0985e344235f5d946f71ac543f1b1565f2cdbc9a2aaee8adf55b",
+    strip_prefix = "rules_python-0.26.0",
+    url = "https://github.com/bazelbuild/rules_python/releases/download/0.26.0/rules_python-0.26.0.tar.gz",
+)
+
 http_archive(
     name = "org_tensorflow",
     sha256 = "ce357fd0728f0d1b0831d1653f475591662ec5bca736a94ff789e6b1944df19f",
@@ -14,6 +33,29 @@ http_archive(
 )
 
 # Copied from `@org_tensorflow//:WORKSPACE`.
+load(
+    "@rules_python//python:repositories.bzl",
+    "py_repositories",
+    "python_register_toolchains",
+)
+py_repositories()
+
+load(
+    "@org_tensorflow//tensorflow/tools/toolchains/python:python_repo.bzl",
+    "python_repository",
+)
+python_repository(name = "python_version_repo")
+load("@python_version_repo//:py_version.bzl", "HERMETIC_PYTHON_VERSION")
+
+# TF workspace scripts below requires `@python` toolchains repo.
+# Toolchain setup here is to please the TF workspace scripts,
+# and we do not use this Python version to build pip packages.
+python_register_toolchains(
+    name = "python",
+    ignore_root_user_error = True,
+    python_version = HERMETIC_PYTHON_VERSION,
+)
+
 load("@org_tensorflow//tensorflow:workspace3.bzl", "tf_workspace3")
 tf_workspace3()