Use legacy optimizer in TPU embedding.

PiperOrigin-RevId: 475892987

Author: chenmoneygithub

site/en/guide/migrate/tpu_embedding.ipynb

@@ -476,7 +476,10 @@
       "source": [
         "strategy = tf.distribute.TPUStrategy(cluster_resolver)\n",
         "with strategy.scope():\n",
-        "  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)\n",
+        "  if hasattr(tf.keras.optimizers, \"legacy\"):\n",
+        "    optimizer = tf.keras.optimizers.legacy.Adagrad(learning_rate=0.05)\n",
+        "  else:\n",
+        "    optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)\n",
         "  dense_input = tf.keras.Input(shape=(2,), dtype=tf.float32, batch_size=global_batch_size)\n",
         "  sparse_input = tf.keras.Input(shape=(), dtype=tf.int32, batch_size=global_batch_size)\n",
         "  embedded_input = tfrs.layers.embedding.TPUEmbedding(\n",
@@ -574,6 +577,7 @@
     "colab": {
       "collapsed_sections": [],
       "name": "tpu_embedding.ipynb",
+      "provenance": [],
       "toc_visible": true
     },
     "kernelspec": {