Removes remaining legacy imports, minor corrections.

Johannes Ballé · copybara-github · commit 0028f42eb58f · 2021-02-03T00:55:50.000-08:00
PiperOrigin-RevId: 355342833
Change-Id: I5050a68312dadedf9b0eb61da541bd77b9f17fed
diff --git a/models/toy_sources/compression_model.py b/models/toy_sources/compression_model.py
@@ -3,7 +3,7 @@
 import abc
 import matplotlib.pyplot as plt
 import numpy as np
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_probability as tfp
 
 
diff --git a/models/toy_sources/ntc.py b/models/toy_sources/ntc.py
@@ -2,7 +2,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_probability as tfp
 import tensorflow_compression as tfc
 from toy_sources import compression_model
diff --git a/models/toy_sources/sawbridge.py b/models/toy_sources/sawbridge.py
@@ -1,6 +1,6 @@
 """Sawbridge process."""
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_probability as tfp
 
 
diff --git a/models/toy_sources/toy_sources.ipynb b/models/toy_sources/toy_sources.ipynb
@@ -58,7 +58,7 @@
         "\n",
         "from absl import logging\n",
         "import numpy as np\n",
-        "import tensorflow.compat.v2 as tf\n",
+        "import tensorflow as tf\n",
         "import tensorflow_probability as tfp\n",
         "import tensorflow_compression as tfc\n",
         "\n",
diff --git a/models/toy_sources/vecvq.py b/models/toy_sources/vecvq.py
@@ -1,6 +1,6 @@
 """Variational entropy-constrained vector quantization."""
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 from toy_sources import compression_model
 
 
diff --git a/tensorflow_compression/all_tests.py b/tensorflow_compression/all_tests.py
@@ -17,7 +17,7 @@
 This is a convenience file to be included in the pip package.
 """
 
-import tensorflow.compat.v1 as tf
+import tensorflow as tf
 
 # pylint: disable=wildcard-import
 from tensorflow_compression.python.distributions.deep_factorized_test import *
diff --git a/tensorflow_compression/python/distributions/deep_factorized.py b/tensorflow_compression/python/distributions/deep_factorized.py
@@ -14,14 +14,16 @@
 # ==============================================================================
 """Deep fully factorized distribution based on cumulative."""
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_probability as tfp
-
 from tensorflow_compression.python.distributions import helpers
 from tensorflow_compression.python.distributions import uniform_noise
 
 
-__all__ = ["DeepFactorized", "NoisyDeepFactorized"]
+__all__ = [
+    "DeepFactorized",
+    "NoisyDeepFactorized",
+]
 
 
 def log_expm1(x):
diff --git a/tensorflow_compression/python/distributions/deep_factorized_test.py b/tensorflow_compression/python/distributions/deep_factorized_test.py
@@ -15,9 +15,8 @@
 """Tests of deep factorized distribution."""
 
 from absl.testing import parameterized
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_probability as tfp
-
 from tensorflow_compression.python.distributions import deep_factorized
 from tensorflow_compression.python.distributions import helpers
 
diff --git a/tensorflow_compression/python/distributions/helpers.py b/tensorflow_compression/python/distributions/helpers.py
@@ -14,7 +14,7 @@
 # ==============================================================================
 """Extensions to `tfp.distributions.Distribution` to handle range coding."""
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 
 
 __all__ = [
diff --git a/tensorflow_compression/python/distributions/helpers_test.py b/tensorflow_compression/python/distributions/helpers_test.py
@@ -14,9 +14,8 @@
 # ==============================================================================
 """Tests of uniform noise adapter distribution."""
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_probability as tfp
-
 from tensorflow_compression.python.distributions import deep_factorized
 from tensorflow_compression.python.distributions import helpers
 
diff --git a/tensorflow_compression/python/distributions/round_adapters.py b/tensorflow_compression/python/distributions/round_adapters.py
@@ -13,19 +13,23 @@
 # limitations under the License.
 # ==============================================================================
 """Distribution adapters for (soft) round functions."""
+
 import tensorflow as tf
 import tensorflow_probability as tfp
-
 from tensorflow_compression.python.distributions import deep_factorized
 from tensorflow_compression.python.distributions import helpers
 from tensorflow_compression.python.distributions import uniform_noise
 from tensorflow_compression.python.ops import soft_round_ops
 
 
 __all__ = [
-    "MonotonicAdapter", "RoundAdapter", "NoisyRoundedNormal",
-    "NoisyRoundedDeepFactorized", "SoftRoundAdapter", "NoisySoftRoundedNormal",
-    "NoisySoftRoundedDeepFactorized"
+    "MonotonicAdapter",
+    "RoundAdapter",
+    "NoisyRoundedNormal",
+    "NoisyRoundedDeepFactorized",
+    "SoftRoundAdapter",
+    "NoisySoftRoundedNormal",
+    "NoisySoftRoundedDeepFactorized",
 ]
 
 
diff --git a/tensorflow_compression/python/distributions/round_adapters_test.py b/tensorflow_compression/python/distributions/round_adapters_test.py
@@ -18,7 +18,6 @@
 import numpy as np
 import tensorflow as tf
 import tensorflow_probability as tfp
-
 from tensorflow_compression.python.distributions import deep_factorized
 from tensorflow_compression.python.distributions import round_adapters
 from tensorflow_compression.python.ops import soft_round_ops
@@ -179,7 +178,7 @@ def test_log_prob_gradient_is_bounded_failcase(self):
           values=[1e6, 1e9])
 
 
-class LocationScaleTest(object):
+class LocationScaleTest:
   """Common tests for noisy location-scale family of distributions."""
 
   def test_can_instantiate_scalar(self):
diff --git a/tensorflow_compression/python/distributions/uniform_noise.py b/tensorflow_compression/python/distributions/uniform_noise.py
@@ -14,9 +14,8 @@
 # ==============================================================================
 """Uniform noise adapter distribution."""
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_probability as tfp
-
 from tensorflow_compression.python.distributions import helpers
 
 
diff --git a/tensorflow_compression/python/distributions/uniform_noise_test.py b/tensorflow_compression/python/distributions/uniform_noise_test.py
@@ -14,13 +14,12 @@
 # ==============================================================================
 """Tests of uniform noise adapter distribution."""
 
-import tensorflow.compat.v2 as tf
-
+import tensorflow as tf
 from tensorflow_compression.python.distributions import helpers
 from tensorflow_compression.python.distributions import uniform_noise
 
 
-class LocationScaleTest(object):
+class LocationScaleTest:
   """Common tests for noisy location-scale family of distributions."""
 
   def test_can_instantiate_scalar(self):
@@ -87,7 +86,7 @@ class NoisyLogisticTest(LocationScaleTest, tf.test.TestCase):
   dist_cls = uniform_noise.NoisyLogistic
 
 
-class MixtureTest(object):
+class MixtureTest:
   """Common tests for noisy mixture distributions."""
 
   def test_can_instantiate_scalar(self):
diff --git a/tensorflow_compression/python/entropy_models/continuous_base.py b/tensorflow_compression/python/entropy_models/continuous_base.py
@@ -16,16 +16,16 @@
 
 import abc
 import functools
-
 from absl import logging
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_probability as tfp
-
 from tensorflow_compression.python.distributions import helpers
 from tensorflow_compression.python.ops import range_coding_ops
 
 
-__all__ = ["ContinuousEntropyModelBase"]
+__all__ = [
+    "ContinuousEntropyModelBase",
+]
 
 
 class ContinuousEntropyModelBase(tf.Module, metaclass=abc.ABCMeta):
diff --git a/tensorflow_compression/python/entropy_models/continuous_batched.py b/tensorflow_compression/python/entropy_models/continuous_batched.py
@@ -15,9 +15,7 @@
 """Batched entropy model for continuous random variables."""
 
 import functools
-
-import tensorflow.compat.v2 as tf
-
+import tensorflow as tf
 from tensorflow_compression.python.distributions import helpers
 from tensorflow_compression.python.entropy_models import continuous_base
 from tensorflow_compression.python.ops import math_ops
diff --git a/tensorflow_compression/python/entropy_models/continuous_batched_test.py b/tensorflow_compression/python/entropy_models/continuous_batched_test.py
@@ -14,9 +14,8 @@
 # ==============================================================================
 """Tests of batched continuous entropy model."""
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_probability as tfp
-
 from tensorflow_compression.python.distributions import uniform_noise
 from tensorflow_compression.python.entropy_models.continuous_batched import ContinuousBatchedEntropyModel
 
@@ -162,7 +161,7 @@ def test_compression_works_in_tf_function(self):
     # create the entropy model only the first time the function is called, and
     # store it for the second time.
 
-    class Compressor(object):
+    class Compressor:
 
       def compress(self, values):
         if not hasattr(self, "em"):
diff --git a/tensorflow_compression/python/entropy_models/continuous_indexed.py b/tensorflow_compression/python/entropy_models/continuous_indexed.py
@@ -15,7 +15,6 @@
 """Indexed entropy model for continuous random variables."""
 
 import tensorflow as tf
-
 from tensorflow_compression.python.distributions import helpers
 from tensorflow_compression.python.entropy_models import continuous_base
 from tensorflow_compression.python.ops import math_ops
diff --git a/tensorflow_compression/python/entropy_models/continuous_indexed_test.py b/tensorflow_compression/python/entropy_models/continuous_indexed_test.py
@@ -14,8 +14,7 @@
 # ==============================================================================
 """Tests of indexed continuous entropy model."""
 
-import tensorflow.compat.v2 as tf
-
+import tensorflow as tf
 from tensorflow_compression.python.distributions import uniform_noise
 from tensorflow_compression.python.entropy_models import continuous_indexed
 
diff --git a/tensorflow_compression/python/entropy_models/universal.py b/tensorflow_compression/python/entropy_models/universal.py
@@ -10,16 +10,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Entropy models which implement Univiersal Quantization."""
-import functools
+"""Entropy models which implement universal quantization."""
 
+import functools
 import tensorflow as tf
-
 from tensorflow_compression.python.entropy_models import continuous_batched
 from tensorflow_compression.python.entropy_models import continuous_indexed
 from tensorflow_compression.python.ops import math_ops
 
-__all__ = ["UniversalBatchedEntropyModel", "UniversalIndexedEntropyModel"]
+
+__all__ = [
+    "UniversalBatchedEntropyModel",
+    "UniversalIndexedEntropyModel",
+]
 
 
 def _add_offset_indexes(indexes, num_noise_levels):
diff --git a/tensorflow_compression/python/entropy_models/universal_test.py b/tensorflow_compression/python/entropy_models/universal_test.py
@@ -12,13 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests of indexed continuous entropy model."""
+"""Tests of universal entropy models."""
 
 from absl.testing import parameterized
 import scipy.stats
 import tensorflow as tf
 import tensorflow_probability as tfp
-
 from tensorflow_compression.python.distributions import deep_factorized
 from tensorflow_compression.python.distributions import uniform_noise
 from tensorflow_compression.python.entropy_models import universal
diff --git a/tensorflow_compression/python/layers/gdn.py b/tensorflow_compression/python/layers/gdn.py
@@ -297,8 +297,8 @@ def alpha(self) -> tf.Tensor:
     if self.alpha_parameter is None:
       raise RuntimeError("alpha is not initialized yet. Call build().")
     if callable(self.alpha_parameter):
-      return self.alpha_parameter()
-    return tf.convert_to_tensor(self.alpha_parameter, dtype=self.dtype)
+      return tf.convert_to_tensor(self.alpha_parameter(), dtype=self.dtype)
+    return self.alpha_parameter
 
   @property
   def beta(self) -> tf.Tensor:
@@ -321,8 +321,8 @@ def epsilon(self) -> tf.Tensor:
     if self.epsilon_parameter is None:
       raise RuntimeError("epsilon is not initialized yet. Call build().")
     if callable(self.epsilon_parameter):
-      return self.epsilon_parameter()
-    return tf.convert_to_tensor(self.epsilon_parameter, dtype=self.dtype)
+      return tf.convert_to_tensor(self.epsilon_parameter(), dtype=self.dtype)
+    return self.epsilon_parameter
 
   @property
   def _channel_axis(self):
diff --git a/tensorflow_compression/python/layers/signal_conv.py b/tensorflow_compression/python/layers/signal_conv.py
@@ -402,7 +402,7 @@ def padding(self, value):
     self._padding = value
 
   @property
-  def extra_pad_end(self) -> Optional[bool]:
+  def extra_pad_end(self) -> bool:
     if self._extra_pad_end is None:
       return self.padding.startswith("same_")
     return self._extra_pad_end
diff --git a/tensorflow_compression/python/util/packed_tensors.py b/tensorflow_compression/python/util/packed_tensors.py
@@ -15,15 +15,15 @@
 """Packed tensors in bit sequences."""
 
 import numpy as np
-import tensorflow.compat.v1 as tf
+import tensorflow as tf
 
 
 __all__ = [
     "PackedTensors",
 ]
 
 
-class PackedTensors(object):
+class PackedTensors:
   """Packed representation of compressed tensors.
 
   This class can pack and unpack several tensor values into a single string. It
@@ -85,6 +85,7 @@ def pack(self, tensors, arrays):
       del self._example.features.feature[chr(i)]
       i += 1
 
+  # TODO(jonycgn): Remove this function once all models are converted.
   def unpack(self, tensors):
     """Unpacks `Tensor` values from this object."""
     # Check tensor dtype first for a more informative error message.
diff --git a/tensorflow_compression/python/util/packed_tensors_test.py b/tensorflow_compression/python/util/packed_tensors_test.py
@@ -15,13 +15,10 @@
 """Tests of PackedTensors class."""
 
 import numpy as np
-import tensorflow.compat.v1 as tf
-
-from tensorflow.python.framework import test_util
+import tensorflow as tf
 from tensorflow_compression.python.util import packed_tensors
 
 
-@test_util.deprecated_graph_mode_only
 class PackedTensorsTest(tf.test.TestCase):
 
   def test_pack_unpack(self):
@@ -30,8 +27,8 @@ def test_pack_unpack(self):
     shape = np.array([1, 3], dtype=np.int32)
     arrays = [string, shape]
 
-    string_t = tf.placeholder(tf.string, [1])
-    shape_t = tf.placeholder(tf.int32, [2])
+    string_t = tf.zeros([1], dtype=tf.string)
+    shape_t = tf.zeros([2], dtype=tf.int32)
     tensors = [string_t, shape_t]
 
     packed = packed_tensors.PackedTensors()
