Some minor fixes.

- Adds logic that prevents reinstalling TF within Colab.
- Cosmetic changes.

PiperOrigin-RevId: 449312782
Change-Id: I90084da39bc6413edcfc7a8584411b0432e085c7

Author: Johannes Ballé

models/hific/colab.ipynb

@@ -31,10 +31,15 @@
       },
       "outputs": [],
       "source": [
-        "!pip install tensorflow-compression\n",
+        "# Installs the latest version of TFC compatible with the installed TF version.\n",
+        "!pip install tensorflow-compression~=$(pip show tensorflow | perl -p -0777 -e 's/.*Version: (\\d\\.\\d).*/\\1.0/sg')\n",
+        "\n",
+        "# Downloads the 'models' directory from Github.\n",
         "![[ -e /tfc ]] || git clone https://github.com/tensorflow/compression /tfc\n",
         "%cd /tfc/models\n",
-        "import tfci  # Check if tfci.py is available."
+        "\n",
+        "# Checks if tfci.py is available.\n",
+        "import tfci\n"
       ]
     },
     {

models/toy_sources/toy_sources.ipynb

@@ -37,7 +37,10 @@
         "# Run this cell to install the necessary dependencies when running the notebook\n",
         "# directly in a Colaboratory hosted runtime from Github.\n",
         "\n",
-        "!pip install tensorflow-compression\n",
+        "# Installs the latest version of TFC compatible with the installed TF version.\n",
+        "!pip install tensorflow-compression~=$(pip show tensorflow | perl -p -0777 -e 's/.*Version: (\\d\\.\\d).*/\\1.0/sg')\n",
+        "\n",
+        "# Downloads the 'models' directory from Github.\n",
         "![[ -e /tfc ]] || git clone https://github.com/tensorflow/compression /tfc\n",
         "%cd /tfc/models\n"
       ],

tensorflow_compression/python/datasets/__init__.py

@@ -12,6 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Data loaders based on `tf.data.Dataset`."""
+"""Data loaders, based on `tf.data.Dataset`."""
 
 from tensorflow_compression.python.datasets.y4m_dataset import *

tensorflow_compression/python/ops/__init__.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Ops."""
+"""Operations."""
+
 from tensorflow_compression.python.ops.gen_ops import *
 from tensorflow_compression.python.ops.math_ops import *
 from tensorflow_compression.python.ops.padding_ops import *