Project import generated by Copybara.

PiperOrigin-RevId: 247508635
Change-Id: I3e0798513af39a8c906db3aff48a26166edf5857

Author: Googler

.gitignore

@@ -2,3 +2,6 @@
 *.pyc
 *.so
 *~
+
+# Ignore bazel-* symlinks.
+/bazel-*

BUILD

@@ -5,83 +5,33 @@ py_library(
     srcs = [
         "tensorflow_compression/__init__.py",
         "tensorflow_compression/python/__init__.py",
-        "tensorflow_compression/python/layers/__init__.py",
-        "tensorflow_compression/python/layers/entropy_models.py",
-        "tensorflow_compression/python/layers/gdn.py",
-        "tensorflow_compression/python/layers/initializers.py",
-        "tensorflow_compression/python/layers/parameterizers.py",
-        "tensorflow_compression/python/layers/signal_conv.py",
-        "tensorflow_compression/python/ops/__init__.py",
-        "tensorflow_compression/python/ops/math_ops.py",
-        "tensorflow_compression/python/ops/namespace_helper.py",
-        "tensorflow_compression/python/ops/padding_ops.py",
-        "tensorflow_compression/python/ops/range_coding_ops.py",
-        "tensorflow_compression/python/ops/spectral_ops.py",
     ],
-    data = [
-        "//tensorflow_compression/cc:libtensorflow_compression.so",
+    deps = [
+        "//tensorflow_compression/python/layers",
+        "//tensorflow_compression/python/ops",
     ],
     visibility = ["//visibility:public"],
 )
 
+filegroup(
+    name = "pip_src",
+    srcs = [
+        "MANIFEST.in",
+        "setup.py",
+        "tensorflow_compression/python/all_test.py",
+    ],
+)
+
 sh_binary(
     name = "build_pip_pkg",
     srcs = ["build_pip_pkg.sh"],
     data = [
         "LICENSE",
-        "MANIFEST.in",
         "README.md",
-        "setup.py",
+        ":pip_src",
         ":tensorflow_compression",
+        # The following targets are for Python test files.
+        "//tensorflow_compression/python/layers:py_src",
+        "//tensorflow_compression/python/ops:py_src",
     ],
 )
-
-py_test(
-    name = "entropy_models_test",
-    timeout = "long",
-    srcs = ["tensorflow_compression/python/layers/entropy_models_test.py"],
-    deps = [":tensorflow_compression"],
-)
-
-py_test(
-    name = "gdn_test",
-    srcs = ["tensorflow_compression/python/layers/gdn_test.py"],
-    deps = [":tensorflow_compression"],
-)
-
-py_test(
-    name = "parameterizers_test",
-    srcs = ["tensorflow_compression/python/layers/parameterizers_test.py"],
-    deps = [":tensorflow_compression"],
-)
-
-py_test(
-    name = "signal_conv_test",
-    timeout = "long",
-    srcs = ["tensorflow_compression/python/layers/signal_conv_test.py"],
-    deps = [":tensorflow_compression"],
-)
-
-py_test(
-    name = "math_ops_test",
-    srcs = ["tensorflow_compression/python/ops/math_ops_test.py"],
-    deps = [":tensorflow_compression"],
-)
-
-py_test(
-    name = "padding_ops_test",
-    srcs = ["tensorflow_compression/python/ops/padding_ops_test.py"],
-    deps = [":tensorflow_compression"],
-)
-
-py_test(
-    name = "spectral_ops_test",
-    srcs = ["tensorflow_compression/python/ops/spectral_ops_test.py"],
-    deps = [":tensorflow_compression"],
-)
-
-py_test(
-    name = "range_coding_ops_test",
-    srcs = ["tensorflow_compression/python/ops/range_coding_ops_test.py"],
-    deps = [":tensorflow_compression"],
-)

README.md

@@ -1,3 +1,5 @@
+# TensorFlow Compression
+
 This project contains data compression ops and layers for TensorFlow. The
 project website is at
 [tensorflow.github.io/compression](https://tensorflow.github.io/compression).
@@ -20,9 +22,9 @@ For an introduction to lossy image compression with machine learning, take a
 look at @jonycgn's
 [talk on Learned Image Compression](https://www.youtube.com/watch?v=x_q7cZviXkY).
 
-# Quick start
+## Quick start
 
-## Installing release 1.1 (stable)
+### Installing release 1.1 (stable)
 
 Install TensorFlow 1.13 using one of the methods described in the
 [TensorFlow installation instructions](https://www.tensorflow.org/install).
@@ -46,7 +48,7 @@ cd compression-1.1
 for i in tensorflow_compression/python/*/*_test.py; do python $i; done
 ```
 
-## Installing release 1.2b1 (beta)
+### Installing release 1.2b1 (beta)
 
 Set up an environment in which you can install precompiled binary Python
 packages using the `pip` command. Refer to the
@@ -66,7 +68,7 @@ Desktop for Mac, and run the above command inside a container based on the
 [TensorFlow docker image](https://www.tensorflow.org/install/docker) for
 Python 2.7.***
 
-## Using the library
+### Using the library
 
 We recommend importing the library from your Python code as follows:
 
@@ -75,7 +77,7 @@ import tensorflow as tf
 import tensorflow_compression as tfc
 ```
 
-## Using a pre-trained model to compress an image
+### Using a pre-trained model to compress an image
 
 ***Note: you need to have a release >1.1 installed for pre-trained model
 support.***
@@ -106,7 +108,7 @@ will decompress a TFCI file and write a PNG file. By default, an output file
 will be named like the input file, only with the appropriate file extension
 appended (any existing extensions will not be removed).
 
-## Training your own model
+### Training your own model
 
 The
 [examples directory](https://github.com/tensorflow/compression/tree/master/examples)
@@ -159,7 +161,7 @@ python bls2017.py [options] compress original.png compressed.bin
 python bls2017.py [options] decompress compressed.bin reconstruction.png
 ```
 
-# Help & documentation
+## Help & documentation
 
 For usage questions and discussions, please head over to our
 [Google group](https://groups.google.com/forum/#!forum/tensorflow-compression).
@@ -174,7 +176,57 @@ There's also an introduction to our `EntropyBottleneck` class
 and a description of the range coding operators
 [here](https://tensorflow.github.io/compression/docs/range_coding.html).
 
-# Authors
+## Building PIP package
+
+This section describes steps to build PIP package.
+
+Docker image `tensorflow/tensorflow:custom-op` is used for PIP package build.
+Note that this is different from `tensorflow/tensorflow:devel`. To be compatible
+with TensorFlow PIP package, the GCC version must match, but
+`tensorflow/tensorflow:devel` has a different GCC version installed.
+
+Inside a Docker container, Git repo is cloned from GitHub, then run script
+`//:build_pip_pkg`.
+
+```bash
+sudo docker run -v /tmp/tensorflow_compression:/tmp/tensorflow_compression \
+    tensorflow/tensorflow:custom-op \
+    bash -c "git clone https://github.com/tensorflow/compression.git
+    /tensorflow_compression && cd /tensorflow_compression &&
+    bazel run -c opt :build_pip_pkg"
+```
+
+The wheel file is created inside `/tmp/tensorflow_compression`. Optimization
+flags can be passed via `--copt` to the `bazel run` command above. For example,
+`bazel` command above can be `bazel run -c opt --copt=-mavx :build_pip_pkg`.
+
+### Testing PIP package
+
+First install the built wheel file.
+
+```bash
+pip install /tmp/tensorflow_compression/tensorflow_compression-*.whl
+```
+
+And run tests. (Do not run the tests in the workspace directory where
+`WORKSPACE` of `tensorflow_compression` repo lives. In that case, Python
+interpreter attempts to import `tensorflow_compression` packages from the source
+tree rather than from the installed package system directory.)
+
+```bash
+pushd /tmp
+python -m tensorflow_compression.python.all_test
+popd
+```
+
+When done, uninstall the PIP package. The package name is
+`tensorflow-compression` with a hyphen(-).
+
+```bash
+pip uninstall tensorflow-compression
+```
+
+## Authors
 
 Johannes Ballé (github: [jonycgn](https://github.com/jonycgn)), Sung Jin Hwang
 (github: [ssjhv](https://github.com/ssjhv)), and Nick Johnston (github:

WORKSPACE

@@ -1,8 +1,4 @@
-local_repository(
-    name = "tf_custom_op",
-    # Change if tensorflow/custom-op repo is cloned into a different directory.
-    path = "/custom-op",
-)
+workspace(name = "tensorflow_compression")
 
-load("//:workspace.bzl", "tf_compression_workspace")
-tf_compression_workspace()
+load("//:workspace.bzl", "tensorflow_compression_workspace")
+tensorflow_compression_workspace()

build_pip_pkg.sh

@@ -19,38 +19,49 @@
 # ==============================================================================
 
 set -e
-set -x
-
-PIP_FILE_PREFIX="bazel-bin/build_pip_pkg.runfiles/__main__/"
 
+# This script needs some improvements.
+#   - Needs a flag to use 'python' or 'python3' to run setup.py.
+#   - Needs a flag to control shared library file extension.
 function main() {
-  DEST=${1}
-  if [[ -z ${DEST} ]]; then
-    DEST=/tmp
+  if [ -z "${1}" ]; then
+    DEST=/tmp/tensorflow_compression
+  else
+    DEST="${1}"
   fi
 
-  mkdir -p ${DEST}
-  DEST=$(readlink -f "${DEST}")
-  echo "=== destination directory: ${DEST}"
+  mkdir -p "${DEST}"
+  DEST="$(readlink -f "${DEST}")"
 
   TMPDIR=$(mktemp -d -t tmp.XXXXXXXXXX)
-  echo $(date) : "=== Using tmpdir: ${TMPDIR}"
+  trap 'rm -rf "${TMPDIR}"' EXIT
 
-  echo "=== Copying files"
+  echo $(date) : "=== Using tmpdir: ${TMPDIR}"
   PKGDIR="${TMPDIR}/tensorflow_compression"
-  rsync -avm -L --exclude='*_test.py' --exclude='build_pip_pkg*' ${PIP_FILE_PREFIX} "${TMPDIR}"
-  for FILENAME in "LICENSE" "README.md"
-  do
-    mv "${TMPDIR}/${FILENAME}" "${PKGDIR}"
-  done
+
+  echo $(date) : "=== Copying files"
+  rsync -amqL tensorflow_compression/ "${PKGDIR}/"
+  cp MANIFEST.in setup.py "${TMPDIR}"
+  cp LICENSE README.md "${PKGDIR}"
+
+  pushd ${TMPDIR} > /dev/null
+
+  # Check if shared library file is copied. If this fails, it is likely that
+  # this script was run directly from the workspace directory. This should be
+  # run inside bazel-bin/build_pip_pkg.runfiles/tensorflow_compression/, or
+  # simply run using "bazel run".
+  if [ ! -f tensorflow_compression/cc/libtensorflow_compression.so ]; then
+    echo "libtensorflow_compression.so not found. Did you use \"bazel run?\""
+    exit 1
+  fi
 
   echo $(date) : "=== Building wheel"
-  pushd ${TMPDIR}
-  python setup.py bdist_wheel > /dev/null
-  cp dist/*.whl "${DEST}"
-  popd
+  python setup.py bdist_wheel
+
+  echo $(date) : "=== Copying wheel to /tmp/"
+  rsync -a --info=name dist/*.whl "${DEST}"
+  popd > /dev/null
 
-  rm -rf ${TMPDIR}
   echo $(date) : "=== Output wheel file is in: ${DEST}"
 }
 

tensorflow_compression/cc/BUILD

@@ -25,7 +25,6 @@ cc_binary(
         # as, say "org_tensorflow" and replace the deps with
         #   "@org_tensorflow//tensorflow/core:framework",
         #   "@org_tensorflow//tensorflow/core:lib",
-        "@local_config_tf//:libtensorflow_framework",
-        "@local_config_tf//:tf_header_lib",
+        "@tensorflow_pip//:libtensorflow_framework",
     ],
 )

tensorflow_compression/python/all_test.py

@@ -0,0 +1,41 @@
+# Copyright 2019 Google LLC. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Tests for tensorflow_compression/python.
+
+This is a convenience file to be included in PIP package.
+No BUILD entry exists for this file on purpose.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+# pylint: disable=wildcard-import
+from tensorflow_compression.python.layers.entropy_models_test import *
+from tensorflow_compression.python.layers.gdn_test import *
+from tensorflow_compression.python.layers.parameterizers_test import *
+from tensorflow_compression.python.layers.signal_conv_test import *
+
+from tensorflow_compression.python.ops.math_ops_test import *
+from tensorflow_compression.python.ops.padding_ops_test import *
+from tensorflow_compression.python.ops.range_coding_ops_test import *
+from tensorflow_compression.python.ops.spectral_ops_test import *
+# pylint: enable=wildcard-import
+
+
+if __name__ == "__main__":
+  tf.test.main()