Update tf.optimizers -> tf_keras.optimizers in STS example notebook.

PiperOrigin-RevId: 606694059

Author: jburnim

tensorflow_probability/examples/jupyter_notebooks/Structural_Time_Series_Modeling_Case_Studies_Atmospheric_CO2_and_Electricity_Demand.ipynb

@@ -85,7 +85,6 @@
         "#@title Import and set ups{ display-mode: \"form\" }\n",
         "\n",
         "%matplotlib inline\n",
-        "import matplotlib as mpl\n",
         "from matplotlib import pylab as plt\n",
         "import matplotlib.dates as mdates\n",
         "import seaborn as sns\n",
@@ -94,12 +93,10 @@
         "\n",
         "import numpy as np\n",
         "import tensorflow.compat.v2 as tf\n",
+        "import tf_keras\n",
         "import tensorflow_probability as tfp\n",
         "\n",
-        "from tensorflow_probability import distributions as tfd\n",
-        "from tensorflow_probability import sts\n",
-        "\n",
-        "tf.enable_v2_behavior()"
+        "from tensorflow_probability import sts"
       ]
     },
     {
@@ -296,7 +293,6 @@
         "\n",
         "  fig=plt.figure(figsize=(12, 6))\n",
         "  ax = fig.add_subplot(1,1,1)\n",
-        "  num_timesteps = one_step_mean.shape[-1]\n",
         "  ax.plot(dates, observed_time_series, label=\"observed time series\", color=c1)\n",
         "  ax.plot(dates, one_step_mean, label=\"one-step prediction\", color=c2)\n",
         "  ax.fill_between(dates,\n",
@@ -504,7 +500,7 @@
         "    target_log_prob_fn=co2_model.joint_distribution(\n",
         "        observed_time_series=co2_by_month_training_data).log_prob,\n",
         "    surrogate_posterior=variational_posteriors,\n",
-        "    optimizer=tf.optimizers.Adam(learning_rate=0.1),\n",
+        "    optimizer=tf_keras.optimizers.Adam(learning_rate=0.1),\n",
         "    num_steps=num_variational_steps,\n",
         "    jit_compile=True)\n",
         "\n",
@@ -918,7 +914,7 @@
         "    target_log_prob_fn=demand_model.joint_distribution(\n",
         "        observed_time_series=demand_training_data).log_prob,\n",
         "    surrogate_posterior=variational_posteriors,\n",
-        "    optimizer=tf.optimizers.Adam(learning_rate=0.1),\n",
+        "    optimizer=tf_keras.optimizers.Adam(learning_rate=0.1),\n",
         "    num_steps=num_variational_steps,\n",
         "    jit_compile=True)\n",
         "plt.plot(elbo_loss_curve)\n",
@@ -1245,6 +1241,7 @@
   "metadata": {
     "colab": {
       "collapsed_sections": [
+        "uiR4-VOt9NFX",
         "5BVYddeJg-An"
       ],
       "name": "Structural Time Series Modeling Case Studies Atmospheric CO2 and Electricity Demand",