Simplified imports and fixed gradients for likelihood lower bound.

Author: Johannes Ballé

tensorflow_compression/__init__.py

@@ -23,17 +23,17 @@
 
 from tensorflow.python.util.all_util import remove_undocumented
 
-# pylint: disable=wildcard-import,g-bad-import-order
-from tensorflow.contrib.coder.python.ops.coder_ops import *
+# pylint: disable=wildcard-import
 from tensorflow_compression.python.layers.entropy_models import *
 from tensorflow_compression.python.layers.gdn import *
 from tensorflow_compression.python.layers.initializers import *
 from tensorflow_compression.python.layers.parameterizers import *
 from tensorflow_compression.python.layers.signal import *
+from tensorflow_compression.python.ops.coder_ops import *
 from tensorflow_compression.python.ops.math_ops import *
 from tensorflow_compression.python.ops.padding_ops import *
 from tensorflow_compression.python.ops.spectral_ops import *
-# pylint: enable=wildcard-import,g-bad-import-order
+# pylint: enable=wildcard-import
 
 remove_undocumented(__name__, [
     "EntropyBottleneck", "GDN", "IdentityInitializer", "Parameterizer",

tensorflow_compression/python/layers/entropy_models.py

@@ -23,8 +23,6 @@
 
 import numpy as np
 
-from tensorflow.contrib.coder.python.ops import coder_ops
-
 from tensorflow.python.eager import context
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
@@ -40,6 +38,9 @@
 from tensorflow.python.ops import state_ops
 from tensorflow.python.summary import summary
 
+from tensorflow_compression.python.ops import coder_ops
+from tensorflow_compression.python.ops import math_ops as tfc_math_ops
+
 
 class EntropyBottleneck(base_layer.Layer):
   """Entropy bottleneck layer.
@@ -436,8 +437,7 @@ def call(self, inputs, training):
     if self.likelihood_bound > 0:
       likelihood_bound = constant_op.constant(
           self.likelihood_bound, dtype=self.dtype)
-      # TODO(jonycgn): Override gradients.
-      likelihood = math_ops.maximum(likelihood, likelihood_bound)
+      likelihood = tfc_math_ops.lower_bound(likelihood, likelihood_bound)
 
     # Convert back to input tensor shape.
     order = list(range(1, ndim))