Fix keras importing failures at the colab.

PiperOrigin-RevId: 607162489

Author: abattery

tensorflow_model_optimization/g3doc/guide/clustering/clustering_example.ipynb

@@ -120,7 +120,7 @@
       "outputs": [],
       "source": [
         "import tensorflow as tf\n",
-        "import tf_keras as keras\n",
+        "from tensorflow_model_optimization.python.core.keras.compat import keras\n",
         "\n",
         "import numpy as np\n",
         "import tempfile\n",

tensorflow_model_optimization/g3doc/guide/pruning/pruning_with_keras.ipynb

@@ -123,7 +123,7 @@
         "import tensorflow as tf\n",
         "import numpy as np\n",
         "\n",
-        "import tf_keras as keras\n",
+        "from tensorflow_model_optimization.python.core.keras.compat import keras\n",
         "\n",
         "%load_ext tensorboard"
       ]

tensorflow_model_optimization/g3doc/guide/quantization/training_example.ipynb

@@ -125,7 +125,7 @@
         "\n",
         "import tensorflow as tf\n",
         "\n",
-        "from tensorflow import keras"
+        "from tensorflow_model_optimization.python.core.keras.compat import keras"
       ]
     },
     {
@@ -146,7 +146,7 @@
       "outputs": [],
       "source": [
         "# Load MNIST dataset\n",
-        "mnist = keras.datasets.mnist\n",
+        "mnist = tf.keras.datasets.mnist\n",
         "(train_images, train_labels), (test_images, test_labels) = mnist.load_data()\n",
         "\n",
         "# Normalize the input image so that each pixel value is between 0 to 1.\n",
@@ -216,7 +216,6 @@
       "outputs": [],
       "source": [
         "import tensorflow_model_optimization as tfmot\n",
-        "import tf_keras as keras\n",
         "\n",
         "quantize_model = tfmot.quantization.keras.quantize_model\n",
         "\n",