Import coder_ops from tf.contrib everywhere, plus some cosmetics.

Author: Johannes Ballé

README.md

@@ -46,8 +46,8 @@ The layer implements a flexible probability density model to estimate entropy,
 which is described in the appendix of the paper (please cite the paper if you
 use this code for scientific work):
 
+> J. Ballé, D. Minnen, S. Singh, S. J. Hwang, N. Johnston:
 > "Variational image compression with a scale hyperprior"
-> Johannes Ballé, David Minnen, Saurabh Singh, Sung Jin Hwang, Nick Johnston
 > https://arxiv.org/abs/1802.01436
 
 The layer assumes that the input tensor is at least 2D, with a batch dimension

compression/__init__.py

@@ -21,8 +21,14 @@
 
 # pylint: disable=wildcard-import
 from compression.python.layers.entropybottleneck import *
-from compression.python.ops.coder_ops import *
+# TODO(jonycgn,nmjohn): revert to importing from here once we figure out why the
+# nightlies don't work.
+# from compression.python.ops.coder_ops import *
+from tensorflow.contrib.coder.python.ops.coder_ops import *
 # pylint: enable=wildcard-import
 
 from tensorflow.python.util.all_util import remove_undocumented
-remove_undocumented(__name__)
+remove_undocumented(__name__, [
+    "EntropyBottleneck",
+    "pmf_to_quantized_cdf", "range_decode", "range_encode",
+])

compression/compile.sh

@@ -22,7 +22,14 @@ TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.ge
 OPTIONAL_MAC_FLAG=''
 
 if [ "$uname" == "Darwin" ]; then
-   OPTIONAL_MAC_FLAG='-undefined dynamic_lookup'
+    OPTIONAL_MAC_FLAG='-undefined dynamic_lookup'
 fi
 
-g++ -std=c++11 -shared kernels/pmf_to_cdf_op.cc kernels/range_coder.cc kernels/range_coder_ops.cc kernels/range_coder_ops_util.cc ops/coder_ops.cc -o _coder_ops.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2 ${OPTIONAL_MAC_FLAG}
+g++ -std=c++11 -shared \
+    kernels/pmf_to_cdf_op.cc \
+    kernels/range_coder.cc \
+    kernels/range_coder_ops.cc \
+    kernels/range_coder_ops_util.cc \
+    ops/coder_ops.cc \
+    -o _coder_ops.so -fPIC \
+    ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2 ${OPTIONAL_MAC_FLAG}

compression/python/ops/coder_ops.py

@@ -20,11 +20,11 @@
 from __future__ import print_function
 
 import os
-
 import tensorflow as tf
 
 
-_coder_ops = tf.load_op_library(os.path.join(tf.resource_loader.get_data_files_path(), "../../_coder_ops.so"))
+_coder_ops = tf.load_op_library(os.path.join(
+    tf.resource_loader.get_data_files_path(), "../../_coder_ops.so"))
 
 # TODO(nmjohn): Find a way to do the below mapping in this module automatically.
 pmf_to_quantized_cdf = _coder_ops.pmf_to_quantized_cdf