Make new Model methods public.

PiperOrigin-RevId: 297727907
Change-Id: I8734601149976bb46c5c37988d53208806ee76b7

tensorflow/python/keras/engine/training.py

@@ -173,6 +173,12 @@ class Model(network.Network, version_utils.ModelVersionSelector):
     self._training_state = None
     self.history = None
 
+    # These objects are used in the default `Model.compile`. They are not
+    # guaranteed to be set after `Model.compile` is called, as users can
+    # override compile with custom logic.
+    self.compiled_loss = None
+    self.compiled_metrics = None
+
   def get_weights(self):
     """Retrieves the weights of the model.
 
@@ -349,9 +355,12 @@ class Model(network.Network, version_utils.ModelVersionSelector):
     """Returns the model's metrics added using `compile`, `add_metric` APIs."""
     metrics = []
     if self._is_compiled:
-      # TODO(omalleyt): Track `CompiledLoss` and `CompiledMetrics` objects
+      # TODO(omalleyt): Track `LossesContainer` and `MetricsContainer` objects
       # so that attr names are not load-bearing.
-      metrics = self.compiled_loss.metrics + self.compiled_metrics.metrics
+      if self.compiled_loss is not None:
+        metrics += self.compiled_loss.metrics
+      if self.compiled_metrics is not None:
+        metrics += self.compiled_metrics.metrics
 
     all_layers = self._gather_unique_layers()
     for l in all_layers:
@@ -414,7 +423,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
   def run_eagerly(self, value):
     self._run_eagerly = value
 
-  def _train_step(self, data):
+  def train_step(self, data):
     """The logic for one training step.
 
     This method can be overridden to support custom training logic.
@@ -462,7 +471,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
     self.compiled_metrics.update_state(y, y_pred, sample_weight)
     return {m.name: m.result() for m in self.metrics}
 
-  def _make_train_function(self):
+  def make_train_function(self):
     """Creates a function that executes one step of training.
 
     This method can be overridden to support custom training logic.
@@ -488,7 +497,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
     def train_function(iterator):
       data = next(iterator)
       outputs = self.distribute_strategy.experimental_run_v2(
-          self._train_step, args=(data,))
+          self.train_step, args=(data,))
       outputs = reduce_per_replica(
           outputs, self.distribute_strategy, reduction='first')
       return outputs
@@ -747,7 +756,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
           steps=data_handler._steps)  # pylint: disable=protected-access
 
       self.stop_training = False
-      train_function = self._make_train_function()
+      train_function = self.make_train_function()
       callbacks.on_train_begin()
       # Handle fault-tolerance for multi-worker.
       # TODO(omalleyt): Fix the ordering issues that mean this has to
@@ -799,7 +808,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
       callbacks.on_train_end()
       return self.history
 
-  def _test_step(self, data):
+  def test_step(self, data):
     """The logic for one evaluation step.
 
     This method can be overridden to support custom evaluation logic.
@@ -833,7 +842,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
     self.compiled_metrics.update_state(y, y_pred, sample_weight)
     return {m.name: m.result() for m in self.metrics}
 
-  def _make_test_function(self):
+  def make_test_function(self):
     """Creates a function that executes one step of evaluation.
 
     This method can be overridden to support custom evaluation logic.
@@ -858,7 +867,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
     def test_function(iterator):
       data = next(iterator)
       outputs = self.distribute_strategy.experimental_run_v2(
-          self._test_step, args=(data,))
+          self.test_step, args=(data,))
       outputs = reduce_per_replica(
           outputs, self.distribute_strategy, reduction='first')
       return outputs
@@ -986,7 +995,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
             epochs=1,
             steps=data_handler._steps)  # pylint: disable=protected-access
 
-      test_function = self._make_test_function()
+      test_function = self.make_test_function()
       callbacks.on_test_begin()
       for _, iterator in data_handler.enumerate_epochs():  # Single epoch.
         self.reset_metrics()
@@ -1012,7 +1021,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
           return results[0]
         return results
 
-  def _predict_step(self, data):
+  def predict_step(self, data):
     """The logic for one inference step.
 
     This method can be overridden to support custom inference logic.
@@ -1036,7 +1045,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
     x, _, _ = data_adapter.unpack_x_y_sample_weight(data)
     return self(x, training=False)
 
-  def _make_predict_function(self):
+  def make_predict_function(self):
     """Creates a function that executes one step of inference.
 
     This method can be overridden to support custom inference logic.
@@ -1060,7 +1069,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
     def predict_function(iterator):
       data = next(iterator)
       outputs = self.distribute_strategy.experimental_run_v2(
-          self._predict_step, args=(data,))
+          self.predict_step, args=(data,))
       outputs = reduce_per_replica(
           outputs, self.distribute_strategy, reduction='concat')
       return outputs
@@ -1173,7 +1182,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
           epochs=1,
           steps=data_handler._steps)  # pylint: disable=protected-access
 
-      predict_function = self._make_predict_function()
+      predict_function = self.make_predict_function()
       callbacks.on_predict_begin()
       for _, iterator in data_handler.enumerate_epochs():  # Single epoch.
         with data_handler.catch_stop_iteration():
@@ -1254,7 +1263,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
       iterator = data_adapter.single_batch_iterator(self.distribute_strategy, x,
                                                     y, sample_weight,
                                                     class_weight)
-      train_function = self._make_train_function()
+      train_function = self.make_train_function()
       logs = train_function(iterator)
 
     if reset_metrics:
@@ -1312,7 +1321,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
     with self.distribute_strategy.scope():
       iterator = data_adapter.single_batch_iterator(self.distribute_strategy, x,
                                                     y, sample_weight)
-      test_function = self._make_test_function()
+      test_function = self.make_test_function()
       logs = test_function(iterator)
 
     if reset_metrics:
@@ -1344,7 +1353,7 @@ class Model(network.Network, version_utils.ModelVersionSelector):
     self._check_call_args('predict_on_batch')
     with self.distribute_strategy.scope():
       iterator = data_adapter.single_batch_iterator(self.distribute_strategy, x)
-      predict_function = self._make_predict_function()
+      predict_function = self.make_predict_function()
       outputs = predict_function(iterator)
     return to_numpy(outputs)
 

tensorflow/python/keras/optimizer_v2/optimizer_v2_test.py

@@ -689,7 +689,8 @@ class OptimizersCompatibilityTest(keras_parameterized.TestCase):
           loss='categorical_crossentropy',
           metrics=[],
           run_eagerly=testing_utils.should_run_eagerly())
-      model_v2._make_train_function()
+      if not ops.executing_eagerly_outside_functions():
+        model_v2._make_train_function()
       if test_weights:
         opt_v2.set_weights(opt_v1.get_weights())
 

tensorflow/python/keras/premade/wide_deep.py

@@ -19,7 +19,6 @@ from __future__ import division
 from __future__ import print_function
 
 from tensorflow.python.eager import backprop
-from tensorflow.python.framework import ops
 from tensorflow.python.keras import activations
 from tensorflow.python.keras import backend as K
 from tensorflow.python.keras import layers as layer_module
@@ -110,7 +109,7 @@ class WideDeepModel(keras_training.Model):
     return output
 
   # This does not support gradient scaling and LossScaleOptimizer.
-  def _train_step(self, data):
+  def train_step(self, data):
     x, y, sample_weight = data_adapter.unpack_x_y_sample_weight(data)
     x, y, sample_weight = data_adapter.expand_1d((x, y, sample_weight))
 
@@ -137,9 +136,6 @@ class WideDeepModel(keras_training.Model):
     return {m.name: m.result() for m in self.metrics}
 
   def _make_train_function(self):
-    if ops.executing_eagerly_outside_functions():
-      return super(WideDeepModel, self)._make_train_function()
-
     # Only needed for graph mode and model_to_estimator.
     has_recompiled = self._recompile_weights_loss_and_weighted_metrics()
     self._check_trainable_weights_consistency()

tensorflow/python/keras/saving/hdf5_format.py

@@ -25,6 +25,7 @@ import os
 import numpy as np
 from six.moves import zip  # pylint: disable=redefined-builtin
 
+from tensorflow.python.framework import ops
 from tensorflow.python.keras import backend as K
 from tensorflow.python.keras import optimizers
 from tensorflow.python.keras.saving import model_config as model_config_lib
@@ -191,7 +192,8 @@ def load_model_from_hdf5(filepath, custom_objects=None, compile=True):  # pylint
         # with data to _make_train_function() and so can't load optimizer
         # weights.
         if model._is_graph_network:  # pylint: disable=protected-access
-          model._make_train_function()
+          if not ops.executing_eagerly_outside_functions():
+            model._make_train_function()
           optimizer_weight_values = load_optimizer_weights_from_hdf5_group(f)
           try:
             model.optimizer.set_weights(optimizer_weight_values)

tensorflow/python/keras/saving/hdf5_format_test.py

@@ -577,8 +577,8 @@ class TestWholeModelSaving(test.TestCase, parameterized.TestCase):
       model.add(keras.layers.Dense(2, input_shape=(3,)))
       model.add(keras.layers.Dense(3))
       model.compile(loss='mse', optimizer='sgd', metrics=['acc'])
-      model._make_train_function()
+      if not ops.executing_eagerly_outside_functions():
-
+        model._make_train_function()
       keras.models.save_model(model, saved_model_dir, save_format=save_format)
       model = keras.models.load_model(saved_model_dir)
 
@@ -964,7 +964,8 @@ class TestWeightSavingAndLoadingTFFormat(test.TestCase):
       model.add(keras.layers.Dense(2, input_shape=(3,)))
       model.add(keras.layers.Dense(3))
       model.compile(loss='mse', optimizer=optimizers.Adam(), metrics=['acc'])
-      model._make_train_function()
+      if not ops.executing_eagerly_outside_functions():
+        model._make_train_function()
       temp_dir = self.get_temp_dir()
       prefix = os.path.join(temp_dir, 'ckpt')
       with test.mock.patch.object(logging, 'warning') as mock_log:

tensorflow/tools/api/golden/v1/tensorflow.keras.-model.pbtxt

@@ -257,6 +257,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "predict"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\', \'callbacks\', \'max_queue_size\', \'workers\', \'use_multiprocessing\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\', \'None\', \'10\', \'1\', \'False\'], "
@@ -269,6 +281,10 @@ tf_class {
     name: "predict_on_batch"
     argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -297,6 +313,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -309,6 +329,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v1/tensorflow.keras.-sequential.pbtxt

@@ -262,6 +262,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "pop"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -286,6 +298,10 @@ tf_class {
     name: "predict_proba"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\'], varargs=None, keywords=None, defaults=[\'32\', \'0\'], "
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -314,6 +330,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -326,6 +346,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-linear-model.pbtxt

@@ -258,6 +258,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "predict"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\', \'callbacks\', \'max_queue_size\', \'workers\', \'use_multiprocessing\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\', \'None\', \'10\', \'1\', \'False\'], "
@@ -270,6 +282,10 @@ tf_class {
     name: "predict_on_batch"
     argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -298,6 +314,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -310,6 +330,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-wide-deep-model.pbtxt

@@ -258,6 +258,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "predict"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\', \'callbacks\', \'max_queue_size\', \'workers\', \'use_multiprocessing\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\', \'None\', \'10\', \'1\', \'False\'], "
@@ -270,6 +282,10 @@ tf_class {
     name: "predict_on_batch"
     argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -298,6 +314,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -310,6 +330,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v1/tensorflow.keras.models.-model.pbtxt

@@ -257,6 +257,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "predict"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\', \'callbacks\', \'max_queue_size\', \'workers\', \'use_multiprocessing\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\', \'None\', \'10\', \'1\', \'False\'], "
@@ -269,6 +281,10 @@ tf_class {
     name: "predict_on_batch"
     argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -297,6 +313,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -309,6 +329,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v1/tensorflow.keras.models.-sequential.pbtxt

@@ -262,6 +262,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "pop"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -286,6 +298,10 @@ tf_class {
     name: "predict_proba"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\'], varargs=None, keywords=None, defaults=[\'32\', \'0\'], "
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -314,6 +330,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -326,6 +346,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v2/tensorflow.keras.-model.pbtxt

@@ -257,6 +257,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "predict"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\', \'callbacks\', \'max_queue_size\', \'workers\', \'use_multiprocessing\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\', \'None\', \'10\', \'1\', \'False\'], "
@@ -269,6 +281,10 @@ tf_class {
     name: "predict_on_batch"
     argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -297,6 +313,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -309,6 +329,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v2/tensorflow.keras.-sequential.pbtxt

@@ -262,6 +262,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "pop"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -286,6 +298,10 @@ tf_class {
     name: "predict_proba"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\'], varargs=None, keywords=None, defaults=[\'32\', \'0\'], "
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -314,6 +330,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -326,6 +346,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-linear-model.pbtxt

@@ -258,6 +258,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "predict"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\', \'callbacks\', \'max_queue_size\', \'workers\', \'use_multiprocessing\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\', \'None\', \'10\', \'1\', \'False\'], "
@@ -270,6 +282,10 @@ tf_class {
     name: "predict_on_batch"
     argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -298,6 +314,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -310,6 +330,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-wide-deep-model.pbtxt

@@ -258,6 +258,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "predict"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\', \'callbacks\', \'max_queue_size\', \'workers\', \'use_multiprocessing\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\', \'None\', \'10\', \'1\', \'False\'], "
@@ -270,6 +282,10 @@ tf_class {
     name: "predict_on_batch"
     argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -298,6 +314,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -310,6 +330,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v2/tensorflow.keras.models.-model.pbtxt

@@ -257,6 +257,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "predict"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\', \'callbacks\', \'max_queue_size\', \'workers\', \'use_multiprocessing\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\', \'None\', \'10\', \'1\', \'False\'], "
@@ -269,6 +281,10 @@ tf_class {
     name: "predict_on_batch"
     argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -297,6 +313,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -309,6 +329,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v2/tensorflow.keras.models.-sequential.pbtxt

@@ -262,6 +262,18 @@ tf_class {
     name: "load_weights"
     argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
   }
+  member_method {
+    name: "make_predict_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_test_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "make_train_function"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "pop"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -286,6 +298,10 @@ tf_class {
     name: "predict_proba"
     argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\'], varargs=None, keywords=None, defaults=[\'32\', \'0\'], "
   }
+  member_method {
+    name: "predict_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "reset_metrics"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
@@ -314,6 +330,10 @@ tf_class {
     name: "test_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "test_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "to_json"
     argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
@@ -326,6 +346,10 @@ tf_class {
     name: "train_on_batch"
     argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\', \'reset_metrics\', \'return_dict\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'False\'], "
   }
+  member_method {
+    name: "train_step"
+    argspec: "args=[\'self\', \'data\'], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "with_name_scope"
     argspec: "args=[\'cls\', \'method\'], varargs=None, keywords=None, defaults=None"