Adds missing __all__ to soft-round layers.

PiperOrigin-RevId: 362230048
Change-Id: Ice9a8332701d207baf9b4ec4a96b5ff975a75668

Author: Johannes Ballé

tensorflow_compression/python/layers/soft_round.py

@@ -13,11 +13,18 @@
 # limitations under the License.
 # ==============================================================================
 """Layers for soft rounding."""
-import tensorflow as tf
 
+import tensorflow as tf
 from tensorflow_compression.python.ops import soft_round_ops
 
 
+__all__ = [
+    "Round",
+    "SoftRound",
+    "SoftRoundConditionalMean",
+]
+
+
 class Round(tf.keras.layers.Layer):
   """Applies rounding."""
 

tools/generate_docs.py

@@ -30,8 +30,7 @@ def main(_):
       py_modules=[("tfc", tfc)],
       base_dir=os.path.dirname(tfc.__file__),
       private_map={
-          # soft_round_ops module seems to trigger a bug.
-          "tfc": ["python", "soft_round_ops"],
+          "tfc": ["python"],
       },
       code_url_prefix="https://github.com/tensorflow/compression/tree/master/"
                       "tensorflow_compression",