Merge commit for internal changes

tensorflow · Oct 20, 2022 · 88f7306 · 88f7306
2 parents a1759b3 + 461b6d5
commit 88f7306
Show file tree

Hide file tree

Showing 21 changed files with 265 additions and 257 deletions.
diff --git a/docs/tutorials/canned_estimators.ipynb b/docs/tutorials/canned_estimators.ipynb
@@ -450,7 +450,7 @@
         "    feature_columns=feature_columns[:5],\n",
         "    model_config=model_config,\n",
         "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
-        "    optimizer=tf.keras.optimizers.Adam(LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42))\n",
         "estimator.train(input_fn=train_input_fn)\n",
         "results = estimator.evaluate(input_fn=test_input_fn)\n",
@@ -497,7 +497,7 @@
         "    feature_columns=feature_columns[:5],\n",
         "    model_config=model_config,\n",
         "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
-        "    optimizer=tf.keras.optimizers.Adam(LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42))\n",
         "estimator.train(input_fn=train_input_fn)\n",
         "results = estimator.evaluate(input_fn=test_input_fn)\n",
@@ -549,7 +549,7 @@
         "    feature_columns=feature_columns,\n",
         "    model_config=model_config,\n",
         "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
-        "    optimizer=tf.keras.optimizers.Adam(LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42))\n",
         "estimator.train(input_fn=train_input_fn)\n",
         "results = estimator.evaluate(input_fn=test_input_fn)\n",
@@ -600,7 +600,7 @@
         "    feature_columns=feature_columns,\n",
         "    model_config=model_config,\n",
         "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
-        "    optimizer=tf.keras.optimizers.Adam(LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42))\n",
         "estimator.train(input_fn=train_input_fn)\n",
         "results = estimator.evaluate(input_fn=test_input_fn)\n",
@@ -672,8 +672,8 @@
         "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
         "    # prefitting_input_fn is required to train the prefitting model.\n",
         "    prefitting_input_fn=prefitting_input_fn,\n",
-        "    optimizer=tf.keras.optimizers.Adam(LEARNING_RATE),\n",
-        "    prefitting_optimizer=tf.keras.optimizers.Adam(LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(LEARNING_RATE),\n",
+        "    prefitting_optimizer=tf.keras.optimizers.legacy.Adam(LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42))\n",
         "estimator.train(input_fn=train_input_fn)\n",
         "results = estimator.evaluate(input_fn=test_input_fn)\n",

diff --git a/docs/tutorials/shape_constraints.ipynb b/docs/tutorials/shape_constraints.ipynb
@@ -619,7 +619,7 @@
         "    feature_columns=feature_columns,\n",
         "    # Hyper-params optimized on validation set.\n",
         "    hidden_units=[16, 8, 8],\n",
-        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
         ")\n",
         "dnn_estimator.train(input_fn=train_input_fn)\n",
@@ -696,7 +696,7 @@
         "    feature_columns=feature_columns,\n",
         "    model_config=model_config,\n",
         "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
-        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
         ")\n",
         "tfl_estimator.train(input_fn=train_input_fn)\n",
@@ -789,7 +789,7 @@
         "    feature_columns=feature_columns,\n",
         "    model_config=model_config,\n",
         "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
-        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
         ")\n",
         "tfl_estimator.train(input_fn=train_input_fn)\n",
@@ -855,7 +855,7 @@
         "    feature_columns=feature_columns,\n",
         "    model_config=model_config,\n",
         "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
-        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
         ")\n",
         "tfl_estimator.train(input_fn=train_input_fn)\n",
@@ -955,7 +955,7 @@
         "    feature_columns=feature_columns,\n",
         "    model_config=model_config,\n",
         "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
-        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
         ")\n",
         "tfl_estimator.train(input_fn=train_input_fn)\n",
@@ -1080,7 +1080,7 @@
         "    feature_columns=feature_columns,\n",
         "    model_config=model_config,\n",
         "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
-        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
         ")\n",
         "tfl_estimator.train(input_fn=train_input_fn)\n",
@@ -1181,7 +1181,7 @@
         "    feature_columns=feature_columns,\n",
         "    model_config=model_config,\n",
         "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
-        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=LEARNING_RATE),\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
         ")\n",
         "tfl_estimator.train(input_fn=train_input_fn)\n",

diff --git a/docs/tutorials/shape_constraints_for_ethics.ipynb b/docs/tutorials/shape_constraints_for_ethics.ipynb
@@ -427,7 +427,7 @@
         "      feature_columns=feature_columns,\n",
         "      model_config=model_config,\n",
         "      feature_analysis_input_fn=get_input_fn(input_df=train_df, num_epochs=1),\n",
-        "      optimizer=tf.keras.optimizers.Adam(learning_rate))\n",
+        "      optimizer=tf.keras.optimizers.legacy.Adam(learning_rate))\n",
         "\n",
         "  estimator.train(\n",
         "      input_fn=get_input_fn(\n",
@@ -733,7 +733,7 @@
         "        tf.feature_column.numeric_column(feature) for feature in feature_names\n",
         "    ],\n",
         "    hidden_units=[100, 100],\n",
-        "    optimizer=tf.keras.optimizers.Adam(learning_rate=0.008),\n",
+        "    optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=0.008),\n",
         "    activation_fn=tf.nn.relu)\n",
         "\n",
         "dnn_estimator.train(\n",

diff --git a/examples/canned_estimators_uci_heart.py b/examples/canned_estimators_uci_heart.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Lint as: python3
 """Example usage of TFL canned estimators.
 
 This example trains several TFL canned estimators on the UCI heart dataset.
@@ -239,7 +238,7 @@ def main(_):
       feature_columns=feature_columns,
       model_config=model_config,
       feature_analysis_input_fn=feature_analysis_input_fn,
-      optimizer=tf.keras.optimizers.Adam(FLAGS.learning_rate))
+      optimizer=tf.keras.optimizers.legacy.Adam(FLAGS.learning_rate))
   estimator.train(input_fn=train_input_fn)
   results = estimator.evaluate(input_fn=test_input_fn)
   print('Calibrated linear results: {}'.format(results))
@@ -260,7 +259,7 @@ def main(_):
       feature_columns=feature_columns,
       model_config=model_config,
       feature_analysis_input_fn=feature_analysis_input_fn,
-      optimizer=tf.keras.optimizers.Adam(FLAGS.learning_rate))
+      optimizer=tf.keras.optimizers.legacy.Adam(FLAGS.learning_rate))
   estimator.train(input_fn=train_input_fn)
   results = estimator.evaluate(input_fn=test_input_fn)
   print('Calibrated lattice results: {}'.format(results))
@@ -285,7 +284,7 @@ def main(_):
       feature_columns=feature_columns,
       model_config=model_config,
       feature_analysis_input_fn=feature_analysis_input_fn,
-      optimizer=tf.keras.optimizers.Adam(FLAGS.learning_rate))
+      optimizer=tf.keras.optimizers.legacy.Adam(FLAGS.learning_rate))
   estimator.train(input_fn=train_input_fn)
   results = estimator.evaluate(input_fn=test_input_fn)
   print('Random ensemble results: {}'.format(results))
@@ -316,7 +315,7 @@ def main(_):
       feature_analysis_input_fn=feature_analysis_input_fn,
       # prefitting_input_fn is required to train the prefitting model.
       prefitting_input_fn=prefitting_input_fn,
-      optimizer=tf.keras.optimizers.Adam(FLAGS.learning_rate))
+      optimizer=tf.keras.optimizers.legacy.Adam(FLAGS.learning_rate))
   estimator.train(input_fn=train_input_fn)
   results = estimator.evaluate(input_fn=test_input_fn)
   print('Crystals ensemble results: {}'.format(results))

diff --git a/examples/custom_estimators_uci_heart.py b/examples/custom_estimators_uci_heart.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Lint as: python3
 """Example usage of TFL layers in custom estimators.
 
 This example trains a TFL custom estimators on the UCI heart dataset.

diff --git a/examples/keras_functional_uci_heart.py b/examples/keras_functional_uci_heart.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Lint as: python3
 """Example usage of TFL within Keras Functional API.
 
 This example builds and trains a calibrated lattice model for the UCI heart

diff --git a/examples/keras_sequential_uci_heart.py b/examples/keras_sequential_uci_heart.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Lint as: python3
 """Example usage of TFL within Keras models.
 
 This example builds and trains a calibrated lattice model for the UCI heart

diff --git a/tensorflow_lattice/python/categorical_calibration_test.py b/tensorflow_lattice/python/categorical_calibration_test.py
@@ -33,10 +33,11 @@
 class CategoricalCalibrationLayerTest(parameterized.TestCase, tf.test.TestCase):
 
   def setUp(self):
+    super(CategoricalCalibrationLayerTest, self).setUp()
     self._disable_all = False
     self._loss_eps = 1e-2
     self._loss_diff_eps = 1e-4
-    super(CategoricalCalibrationLayerTest, self).setUp()
+    tf.keras.utils.set_random_seed(42)
 
   def _ResetAllBackends(self):
     keras.backend.clear_session()