Exports clean namespace using __all__ mechanism.

PiperOrigin-RevId: 242486798

Author: jballe

__init__.py

@@ -20,8 +20,6 @@
 
 # Dependency imports
 
-from tensorflow.python.util.all_util import remove_undocumented
-
 # pylint: disable=wildcard-import
 from tensorflow_compression.python.layers.entropy_models import *
 from tensorflow_compression.python.layers.gdn import *
@@ -33,11 +31,3 @@
 from tensorflow_compression.python.ops.range_coding_ops import *
 from tensorflow_compression.python.ops.spectral_ops import *
 # pylint: enable=wildcard-import
-
-remove_undocumented(__name__, [
-    "EntropyBottleneck", "GDN", "IdentityInitializer", "Parameterizer",
-    "StaticParameterizer", "RDFTParameterizer", "NonnegativeParameterizer",
-    "SignalConv1D", "SignalConv2D", "SignalConv3D",
-    "upper_bound", "lower_bound", "same_padding_for_kernel", "irdft_matrix",
-    "pmf_to_quantized_cdf", "range_decode", "range_encode",
-])

python/layers/gdn.py

@@ -31,6 +31,11 @@
 _default_gamma_param = parameterizers.NonnegativeParameterizer()
 
 
+__all__ = [
+    "GDN",
+]
+
+
 class GDN(tf.keras.layers.Layer):
   """Generalized divisive normalization layer.
 

python/layers/initializers.py

@@ -24,6 +24,11 @@
 import tensorflow as tf
 
 
+__all__ = [
+    "IdentityInitializer",
+]
+
+
 class IdentityInitializer(object):
   """Initialize to the identity kernel with the given shape.
 

python/layers/parameterizers.py

@@ -27,6 +27,14 @@
 from tensorflow_compression.python.ops import spectral_ops
 
 
+__all__ = [
+    "Parameterizer",
+    "StaticParameterizer",
+    "RDFTParameterizer",
+    "NonnegativeParameterizer",
+]
+
+
 class Parameterizer(object):
   """Parameterization object (abstract base class).
 

python/layers/signal_conv.py

@@ -29,6 +29,13 @@
 from tensorflow_compression.python.ops import padding_ops
 
 
+__all__ = [
+    "SignalConv1D",
+    "SignalConv2D",
+    "SignalConv3D",
+]
+
+
 class _SignalConv(tf.keras.layers.Layer):
   """{rank}D convolution layer.
 

python/ops/math_ops.py

@@ -24,6 +24,12 @@
 import tensorflow as tf
 
 
+__all__ = [
+    "upper_bound",
+    "lower_bound",
+]
+
+
 @tf.RegisterGradient("IdentityFirstOfTwoInputs")
 def _identity_first_of_two_inputs_grad(op, grad):
   """Gradient for `lower_bound` or `upper_bound` if `gradient == 'identity'`.

python/ops/padding_ops.py

@@ -22,6 +22,11 @@
 # Dependency imports
 
 
+__all__ = [
+    "same_padding_for_kernel",
+]
+
+
 def same_padding_for_kernel(shape, corr, strides_up=None):
   """Determine correct amount of padding for `same` convolution.
 

python/ops/spectral_ops.py

@@ -29,6 +29,11 @@
 _matrix_cache = {}
 
 
+__all__ = [
+    "irdft_matrix",
+]
+
+
 def irdft_matrix(shape, dtype=tf.float32):
   """Matrix for implementing kernel reparameterization with `tf.matmul`.