Docstring/namespace fixes.

PiperOrigin-RevId: 449469688
Change-Id: If188fa57655c93139637f1d051c83447875ffcbb

Author: Johannes Ballé

BUILD

@@ -11,12 +11,29 @@ py_library(
     srcs = ["tensorflow_compression/__init__.py"],
     visibility = ["//visibility:public"],
     deps = [
-        "//tensorflow_compression/python/datasets",
+        "//tensorflow_compression/python/datasets:y4m_dataset",
         "//tensorflow_compression/python/distributions",
+        "//tensorflow_compression/python/distributions:deep_factorized",
+        "//tensorflow_compression/python/distributions:helpers",
+        "//tensorflow_compression/python/distributions:round_adapters",
+        "//tensorflow_compression/python/distributions:uniform_noise",
         "//tensorflow_compression/python/entropy_models",
+        "//tensorflow_compression/python/entropy_models:continuous_batched",
+        "//tensorflow_compression/python/entropy_models:continuous_indexed",
+        "//tensorflow_compression/python/entropy_models:power_law",
+        "//tensorflow_compression/python/entropy_models:universal",
         "//tensorflow_compression/python/layers",
+        "//tensorflow_compression/python/layers:gdn",
+        "//tensorflow_compression/python/layers:initializers",
+        "//tensorflow_compression/python/layers:parameters",
+        "//tensorflow_compression/python/layers:signal_conv",
+        "//tensorflow_compression/python/layers:soft_round",
         "//tensorflow_compression/python/ops",
-        "//tensorflow_compression/python/util",
+        "//tensorflow_compression/python/ops:gen_ops",
+        "//tensorflow_compression/python/ops:math_ops",
+        "//tensorflow_compression/python/ops:padding_ops",
+        "//tensorflow_compression/python/ops:round_ops",
+        "//tensorflow_compression/python/util:packed_tensors",
     ],
 )
 

tensorflow_compression/__init__.py

@@ -20,10 +20,23 @@
 from tensorflow_compression.python import ops
 
 # pylint: disable=wildcard-import
-from tensorflow_compression.python.datasets import *
-from tensorflow_compression.python.distributions import *
-from tensorflow_compression.python.entropy_models import *
-from tensorflow_compression.python.layers import *
-from tensorflow_compression.python.ops import *
-from tensorflow_compression.python.util import *
+from tensorflow_compression.python.datasets.y4m_dataset import *
+from tensorflow_compression.python.distributions.deep_factorized import *
+from tensorflow_compression.python.distributions.helpers import *
+from tensorflow_compression.python.distributions.round_adapters import *
+from tensorflow_compression.python.distributions.uniform_noise import *
+from tensorflow_compression.python.entropy_models.continuous_batched import *
+from tensorflow_compression.python.entropy_models.continuous_indexed import *
+from tensorflow_compression.python.entropy_models.power_law import *
+from tensorflow_compression.python.entropy_models.universal import *
+from tensorflow_compression.python.layers.gdn import *
+from tensorflow_compression.python.layers.initializers import *
+from tensorflow_compression.python.layers.parameters import *
+from tensorflow_compression.python.layers.signal_conv import *
+from tensorflow_compression.python.layers.soft_round import *
+from tensorflow_compression.python.ops.gen_ops import *
+from tensorflow_compression.python.ops.math_ops import *
+from tensorflow_compression.python.ops.padding_ops import *
+from tensorflow_compression.python.ops.round_ops import *
+from tensorflow_compression.python.util.packed_tensors import *
 # pylint: enable=wildcard-import

tensorflow_compression/python/distributions/uniform_noise.py

@@ -200,6 +200,7 @@ class NoisyMixtureSameFamily(tfp.distributions.MixtureSameFamily):
 
   def __init__(self, mixture_distribution, components_distribution,
                name="NoisyMixtureSameFamily"):
+    """Initializer, taking the same arguments as `tfpd.MixtureSameFamily`."""
     super().__init__(
         mixture_distribution=mixture_distribution,
         components_distribution=UniformNoiseAdapter(components_distribution),
@@ -252,20 +253,30 @@ class NoisyNormal(UniformNoiseAdapter):
   """Gaussian distribution with additive i.i.d. uniform noise."""
 
   def __init__(self, name="NoisyNormal", **kwargs):
+    """Initializer, taking the same arguments as `tfpd.Normal`."""
     super().__init__(tfp.distributions.Normal(**kwargs), name=name)
 
 
 class NoisyLogistic(UniformNoiseAdapter):
   """Logistic distribution with additive i.i.d. uniform noise."""
 
   def __init__(self, name="NoisyLogistic", **kwargs):
+    """Initializer, taking the same arguments as `tfpd.Logistic`."""
     super().__init__(tfp.distributions.Logistic(**kwargs), name=name)
 
 
 class NoisyNormalMixture(NoisyMixtureSameFamily):
   """Mixture of normal distributions with additive i.i.d. uniform noise."""
 
   def __init__(self, loc, scale, weight, name="NoisyNormalMixture"):
+    """Initializer.
+
+    Args:
+      loc: Location parameters of `tfpd.Normal` component distributions.
+      scale: Scale parameters of `tfpd.Normal` component distributions.
+      weight: `probs` parameter of `tfpd.Categorical` mixture distribution.
+      name: A name for this distribution.
+    """
     super().__init__(
         mixture_distribution=tfp.distributions.Categorical(probs=weight),
         components_distribution=tfp.distributions.Normal(loc=loc, scale=scale),
@@ -277,6 +288,14 @@ class NoisyLogisticMixture(NoisyMixtureSameFamily):
   """Mixture of logistic distributions with additive i.i.d. uniform noise."""
 
   def __init__(self, loc, scale, weight, name="NoisyLogisticMixture"):
+    """Initializer.
+
+    Args:
+      loc: Location parameters of `tfpd.Logistic` component distributions.
+      scale: Scale parameters of `tfpd.Logistic` component distributions.
+      weight: `probs` parameter of `tfpd.Categorical` mixture distribution.
+      name: A name for this distribution.
+    """
     super().__init__(
         mixture_distribution=tfp.distributions.Categorical(probs=weight),
         components_distribution=tfp.distributions.Logistic(

tensorflow_compression/python/entropy_models/continuous_base.py

@@ -110,15 +110,18 @@ def prior(self):
 
   @prior.deleter
   def prior(self):
+    """Prior distribution, used for deriving range coding tables."""
     self._prior = None
 
   @property
   def cdf(self):
+    """The CDFs used by range coding."""
     self._check_compression()
     return tf.convert_to_tensor(self._cdf)
 
   @property
   def cdf_offset(self):
+    """The CDF offsets used by range coding."""
     self._check_compression()
     return tf.convert_to_tensor(self._cdf_offset)
 

tensorflow_compression/python/entropy_models/continuous_batched.py

@@ -265,10 +265,12 @@ def prior_shape_tensor(self):
 
   @property
   def offset_heuristic(self):
+    """Whether to use heuristic to determine quantization offsets."""
     return self._offset_heuristic
 
   @property
   def quantization_offset(self):
+    """The quantization offset used in `quantize` and `compress`."""
     if self._quantization_offset is not None:
       return tf.convert_to_tensor(self._quantization_offset)
     if self.offset_heuristic and not self.compression:

tensorflow_compression/python/entropy_models/continuous_indexed.py

@@ -39,7 +39,7 @@ class ContinuousIndexedEntropyModel(continuous_base.ContinuousEntropyModelBase):
   A typical workflow looks like this:
 
   - Train a model using an instance of this entropy model as a bottleneck,
-    passing the bottleneck tensor through it. With training=True, the model
+    passing the bottleneck tensor through it. With `training=True`, the model
     computes a differentiable upper bound on the number of bits needed to
     compress the bottleneck tensor.
   - For evaluation, get a closer estimate of the number of compressed bits

tensorflow_compression/python/entropy_models/power_law.py

@@ -115,16 +115,15 @@ def __call__(self, bottleneck):
 
   @tf.Module.with_name_scope
   def penalty(self, bottleneck):
-    """Computes cross-entropy penalty.
+    """Computes penalty encouraging compressibility.
 
     Args:
       bottleneck: `tf.Tensor` containing the data to be compressed. Must have at
         least `self.coding_rank` dimensions.
 
     Returns:
-      Penalty, which has the same shape as `bottleneck` without the
-      `self.coding_rank` innermost dimensions, and corresponds to a cross
-      entropy.
+      Penalty value, which has the same shape as `bottleneck` without the
+      `self.coding_rank` innermost dimensions.
     """
     bottleneck = tf.convert_to_tensor(bottleneck, dtype=self.bottleneck_dtype)
     penalty = tf.math.log((abs(bottleneck) + self.alpha) / self.alpha)

tensorflow_compression/python/layers/signal_conv.py

@@ -120,6 +120,7 @@ class _SignalConv(tf.keras.layers.Layer):
     inputs:  |0 1 2 3 4 5 6 7 8|
     outputs: |    0 . 1 . 2    |
     ```
+
     ```
     inputs:  |0 1 2 3 4 5 6 7|
     outputs: |    0 . 1 .    |
@@ -152,6 +153,7 @@ class _SignalConv(tf.keras.layers.Layer):
     inputs:  |0 1 2 3 4 5 6 7 8|
     outputs: |0 . 1 . 2 . 3 . 4|
     ```
+
     ```
     inputs:  |0 1 2 3 4 5 6 7|
     outputs: |0 . 1 . 2 . 3 .|