Rollback the previous change of feature_column test.

PiperOrigin-RevId: 314933875
Change-Id: Ib83c967083a2a26e02e1b38ae6baa205204f4a97

tensorflow/python/feature_column/BUILD

@@ -23,6 +23,7 @@ py_library(
     srcs_version = "PY2AND3",
     deps = [
         ":utils",
+        "@six_archive//:six",
         "//tensorflow/python:array_ops",
         "//tensorflow/python:check_ops",
         "//tensorflow/python:control_flow_ops",
@@ -48,7 +49,13 @@ py_library(
         "//tensorflow/python:variables",
         "//tensorflow/python/eager:context",
         "//tensorflow/python/keras/engine",
-        "@six_archive//:six",
+        # TODO(scottzhu): Remove metrics after we cleanup the keras internal cyclar dependency.
+        #     //third_party/tensorflow/python/feature_column:feature_column
+        #     //third_party/tensorflow/python/keras/engine:engine
+        # .-> //third_party/tensorflow/python/keras/distribute:distribute
+        # |   //third_party/tensorflow/python/keras:metrics
+        # `-- //third_party/tensorflow/python/keras/distribute:distribute
+        "//tensorflow/python/keras:metrics",
     ],
 )
 
@@ -72,7 +79,6 @@ py_library(
         "//tensorflow/python:init_ops",
         "//tensorflow/python:lookup_ops",
         "//tensorflow/python:math_ops",
-        "//tensorflow/python:nn_ops",
         "//tensorflow/python:parsing_ops",
         "//tensorflow/python:platform",
         "//tensorflow/python:sparse_ops",
@@ -85,11 +91,7 @@ py_library(
         "//tensorflow/python:variable_scope",
         "//tensorflow/python:variables",
         "//tensorflow/python/eager:context",
-        "//tensorflow/python/keras:backend",
         "//tensorflow/python/keras:initializers",
-        "//tensorflow/python/keras/engine",
-        "//tensorflow/python/keras/engine:base_layer",
-        "//tensorflow/python/keras/layers",
         "//tensorflow/python/keras/utils:generic_utils",
         "//tensorflow/python/training/tracking",
         "//tensorflow/python/training/tracking:data_structures",
@@ -236,11 +238,8 @@ py_test(
     deps = [
         ":feature_column_v2",
         "//tensorflow/python:client_testlib",
-        "//tensorflow/python:framework_ops",
         "//tensorflow/python:parsing_ops",
-        "//tensorflow/python:training",
         "//tensorflow/python:util",
-        "//tensorflow/python/keras/layers",
     ],
 )
 
@@ -254,37 +253,3 @@ tf_py_test(
         "@absl_py//absl/testing:parameterized",
     ],
 )
-
-tf_py_test(
-    name = "keras_integration_test",
-    size = "medium",
-    srcs = ["keras_integration_test.py"],
-    python_version = "PY3",
-    shard_count = 4,
-    tags = [
-        "nomac",  # TODO(mihaimaruseac): b/127695564
-        "notsan",
-    ],
-    deps = [
-        ":feature_column_py",
-        "//tensorflow/python:client_testlib",
-        "//tensorflow/python/keras",
-        "//third_party/py/numpy",
-        "@absl_py//absl/testing:parameterized",
-    ],
-)
-
-tf_py_test(
-    name = "save_test",
-    size = "medium",
-    srcs = ["save_test.py"],
-    python_version = "PY3",
-    deps = [
-        ":feature_column_v2",
-        "//tensorflow/python:client_testlib",
-        "//tensorflow/python/keras",
-        "//tensorflow/python/keras:combinations",
-        "//third_party/py/numpy",
-        "@absl_py//absl/testing:parameterized",
-    ],
-)

tensorflow/python/feature_column/save_test.py

@@ -1,133 +0,0 @@
-# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Tests for Keras model saving code."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import parameterized
-import numpy as np
-
-from tensorflow.python import keras
-from tensorflow.python.eager import context
-from tensorflow.python.feature_column import feature_column_lib
-from tensorflow.python.framework import sparse_tensor
-from tensorflow.python.keras import combinations
-from tensorflow.python.keras.saving import model_config
-from tensorflow.python.ops import lookup_ops
-from tensorflow.python.platform import test
-
-
-class TestSaveModel(test.TestCase, parameterized.TestCase):
-
-  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
-  def test_saving_with_dense_features(self):
-    cols = [
-        feature_column_lib.numeric_column('a'),
-        feature_column_lib.indicator_column(
-            feature_column_lib.categorical_column_with_vocabulary_list(
-                'b', ['one', 'two']))
-    ]
-    input_layers = {
-        'a': keras.layers.Input(shape=(1,), name='a'),
-        'b': keras.layers.Input(shape=(1,), name='b', dtype='string')
-    }
-
-    fc_layer = feature_column_lib.DenseFeatures(cols)(input_layers)
-    output = keras.layers.Dense(10)(fc_layer)
-
-    model = keras.models.Model(input_layers, output)
-
-    model.compile(
-        loss=keras.losses.MSE,
-        optimizer='rmsprop',
-        metrics=[keras.metrics.categorical_accuracy])
-
-    config = model.to_json()
-    loaded_model = model_config.model_from_json(config)
-
-    inputs_a = np.arange(10).reshape(10, 1)
-    inputs_b = np.arange(10).reshape(10, 1).astype('str')
-
-    with self.cached_session():
-      # Initialize tables for V1 lookup.
-      if not context.executing_eagerly():
-        self.evaluate(lookup_ops.tables_initializer())
-
-      self.assertLen(loaded_model.predict({'a': inputs_a, 'b': inputs_b}), 10)
-
-  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
-  def test_saving_with_sequence_features(self):
-    cols = [
-        feature_column_lib.sequence_numeric_column('a'),
-        feature_column_lib.indicator_column(
-            feature_column_lib.sequence_categorical_column_with_vocabulary_list(
-                'b', ['one', 'two']))
-    ]
-    input_layers = {
-        'a':
-            keras.layers.Input(shape=(None, 1), sparse=True, name='a'),
-        'b':
-            keras.layers.Input(
-                shape=(None, 1), sparse=True, name='b', dtype='string')
-    }
-
-    fc_layer, _ = feature_column_lib.SequenceFeatures(cols)(input_layers)
-    # TODO(tibell): Figure out the right dtype and apply masking.
-    # sequence_length_mask = array_ops.sequence_mask(sequence_length)
-    # x = keras.layers.GRU(32)(fc_layer, mask=sequence_length_mask)
-    x = keras.layers.GRU(32)(fc_layer)
-    output = keras.layers.Dense(10)(x)
-
-    model = keras.models.Model(input_layers, output)
-
-    model.compile(
-        loss=keras.losses.MSE,
-        optimizer='rmsprop',
-        metrics=[keras.metrics.categorical_accuracy])
-
-    config = model.to_json()
-    loaded_model = model_config.model_from_json(config)
-
-    batch_size = 10
-    timesteps = 1
-
-    values_a = np.arange(10, dtype=np.float32)
-    indices_a = np.zeros((10, 3), dtype=np.int64)
-    indices_a[:, 0] = np.arange(10)
-    inputs_a = sparse_tensor.SparseTensor(indices_a, values_a,
-                                          (batch_size, timesteps, 1))
-
-    values_b = np.zeros(10, dtype=np.str)
-    indices_b = np.zeros((10, 3), dtype=np.int64)
-    indices_b[:, 0] = np.arange(10)
-    inputs_b = sparse_tensor.SparseTensor(indices_b, values_b,
-                                          (batch_size, timesteps, 1))
-
-    with self.cached_session():
-      # Initialize tables for V1 lookup.
-      if not context.executing_eagerly():
-        self.evaluate(lookup_ops.tables_initializer())
-
-      self.assertLen(
-          loaded_model.predict({
-              'a': inputs_a,
-              'b': inputs_b
-          }, steps=1), batch_size)
-
-
-if __name__ == '__main__':
-  test.main()

tensorflow/python/keras/engine/BUILD

@@ -426,6 +426,24 @@ tf_py_test(
     ],
 )
 
+tf_py_test(
+    name = "feature_columns_integration_test",
+    size = "medium",
+    srcs = ["feature_columns_integration_test.py"],
+    python_version = "PY3",
+    tags = [
+        "nomac",  # TODO(mihaimaruseac): b/127695564
+        "notsan",
+    ],
+    deps = [
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python/feature_column:feature_column_py",
+        "//tensorflow/python/keras",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
 tf_py_test(
     name = "training_eager_test",
     size = "medium",

tensorflow/python/keras/engine/feature_columns_integration_test.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests specific to Feature Columns and Keras integration."""
+"""Tests specific to Feature Columns integration."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -21,17 +21,11 @@ from __future__ import print_function
 import numpy as np
 
 from tensorflow.python import keras
-from tensorflow.python import tf2
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.feature_column import feature_column_lib as fc
-from tensorflow.python.feature_column import feature_column_v2
 from tensorflow.python.keras import keras_parameterized
 from tensorflow.python.keras import metrics as metrics_module
 from tensorflow.python.keras import testing_utils
-from tensorflow.python.keras.feature_column import dense_features_v2
-from tensorflow.python.keras.optimizer_v2 import gradient_descent
-from tensorflow.python.keras.premade import linear
-from tensorflow.python.keras.premade import wide_deep
 from tensorflow.python.keras.utils import np_utils
 from tensorflow.python.platform import test
 
@@ -304,108 +298,6 @@ class FeatureColumnsIntegrationTest(keras_parameterized.TestCase):
                   loss=keras.losses.BinaryCrossentropy())
     model.fit(dataset)
 
-  def test_serialization_dense_features(self):
-    dense_feature = fc.DenseFeatures([fc.numeric_column('a')])
-    config = keras.layers.serialize(dense_feature)
-    self.assertEqual(config['class_name'], 'DenseFeatures')
-
-    revived = keras.layers.deserialize(config)
-    if tf2.enabled():
-      self.assertIsInstance(revived, dense_features_v2.DenseFeatures)
-    else:
-      self.assertIsInstance(revived, fc.DenseFeatures)
-      self.assertNotIsInstance(revived, dense_features_v2.DenseFeatures)
-
-  # This test is an example for a regression on categorical inputs, i.e.,
-  # the output is 0.4, 0.6, 0.9 when input is 'alpha', 'beta', 'gamma'
-  # separately.
-  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
-  def test_linear_model_with_feature_column(self):
-    vocab_list = ['alpha', 'beta', 'gamma']
-    vocab_val = [0.4, 0.6, 0.9]
-    data = np.random.choice(vocab_list, size=256)
-    y = np.zeros_like(data, dtype=np.float32)
-    for vocab, val in zip(vocab_list, vocab_val):
-      indices = np.where(data == vocab)
-      y[indices] = val + np.random.uniform(
-          low=-0.01, high=0.01, size=indices[0].shape)
-    cat_column = feature_column_v2.categorical_column_with_vocabulary_list(
-        key='symbol', vocabulary_list=vocab_list)
-    ind_column = feature_column_v2.indicator_column(cat_column)
-    dense_feature_layer = dense_features_v2.DenseFeatures([ind_column])
-    linear_model = linear.LinearModel(
-        use_bias=False, kernel_initializer='zeros')
-    combined = keras.Sequential([dense_feature_layer, linear_model])
-    opt = gradient_descent.SGD(learning_rate=0.1)
-    combined.compile(opt, 'mse', [])
-    combined.fit(x={'symbol': data}, y=y, batch_size=32, epochs=10)
-    self.assertAllClose([[0.4], [0.6], [0.9]],
-                        combined.layers[1].dense_layers[0].kernel.numpy(),
-                        atol=0.01)
-
-  # This test is an example for cases where linear and dnn model accepts
-  # same raw input and same transformed inputs, i.e., the raw input is
-  # categorical, and both linear and dnn model accept one hot encoding.
-  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
-  def test_wide_deep_model_with_single_feature_column(self):
-    vocab_list = ['alpha', 'beta', 'gamma']
-    vocab_val = [0.4, 0.6, 0.9]
-    data = np.random.choice(vocab_list, size=256)
-    y = np.zeros_like(data, dtype=np.float32)
-    for vocab, val in zip(vocab_list, vocab_val):
-      indices = np.where(data == vocab)
-      y[indices] = val + np.random.uniform(
-          low=-0.01, high=0.01, size=indices[0].shape)
-    cat_column = feature_column_v2.categorical_column_with_vocabulary_list(
-        key='symbol', vocabulary_list=vocab_list)
-    ind_column = feature_column_v2.indicator_column(cat_column)
-    dense_feature_layer = dense_features_v2.DenseFeatures([ind_column])
-    linear_model = linear.LinearModel(
-        use_bias=False, kernel_initializer='zeros')
-    dnn_model = keras.Sequential([keras.layers.Dense(units=1)])
-    wide_deep_model = wide_deep.WideDeepModel(linear_model, dnn_model)
-    combined = keras.Sequential([dense_feature_layer, wide_deep_model])
-    opt = gradient_descent.SGD(learning_rate=0.1)
-    combined.compile(
-        opt,
-        'mse', [],
-        run_eagerly=testing_utils.should_run_eagerly())
-    combined.fit(x={'symbol': data}, y=y, batch_size=32, epochs=10)
-
-  # This test is an example for cases where linear and dnn model accepts
-  # same raw input but different transformed inputs, i.e,. the raw input is
-  # categorical, and linear model accepts one hot encoding, while dnn model
-  # accepts embedding encoding.
-  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
-  def test_wide_deep_model_with_two_feature_columns(self):
-    vocab_list = ['alpha', 'beta', 'gamma']
-    vocab_val = [0.4, 0.6, 0.9]
-    data = np.random.choice(vocab_list, size=256)
-    y = np.zeros_like(data, dtype=np.float32)
-    for vocab, val in zip(vocab_list, vocab_val):
-      indices = np.where(data == vocab)
-      y[indices] = val + np.random.uniform(
-          low=-0.01, high=0.01, size=indices[0].shape)
-    cat_column = feature_column_v2.categorical_column_with_vocabulary_list(
-        key='symbol', vocabulary_list=vocab_list)
-    ind_column = feature_column_v2.indicator_column(cat_column)
-    emb_column = feature_column_v2.embedding_column(cat_column, dimension=5)
-    linear_feature_layer = dense_features_v2.DenseFeatures([ind_column])
-    linear_model = linear.LinearModel(
-        use_bias=False, kernel_initializer='zeros')
-    combined_linear = keras.Sequential(
-        [linear_feature_layer, linear_model])
-    dnn_model = keras.Sequential([keras.layers.Dense(units=1)])
-    dnn_feature_layer = dense_features_v2.DenseFeatures([emb_column])
-    combined_dnn = keras.Sequential([dnn_feature_layer, dnn_model])
-    wide_deep_model = wide_deep.WideDeepModel(combined_linear, combined_dnn)
-    opt = gradient_descent.SGD(learning_rate=0.1)
-    wide_deep_model.compile(
-        opt,
-        'mse', [],
-        run_eagerly=testing_utils.should_run_eagerly())
-    wide_deep_model.fit(x={'symbol': data}, y=y, batch_size=32, epochs=10)
-
 
 if __name__ == '__main__':
   test.main()

tensorflow/python/keras/premade/linear_test.py

@@ -22,6 +22,7 @@ import numpy as np
 
 from tensorflow.python.eager import backprop
 from tensorflow.python.eager import context
+from tensorflow.python.feature_column import feature_column_v2 as fc
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import sparse_tensor
@@ -29,7 +30,9 @@ from tensorflow.python.keras import backend
 from tensorflow.python.keras import keras_parameterized
 from tensorflow.python.keras import losses
 from tensorflow.python.keras.engine import input_layer
+from tensorflow.python.keras.engine import sequential
 from tensorflow.python.keras.engine import training
+from tensorflow.python.keras.feature_column import dense_features_v2
 from tensorflow.python.keras.layers import core
 from tensorflow.python.keras.optimizer_v2 import gradient_descent
 from tensorflow.python.keras.premade import linear
@@ -126,6 +129,33 @@ class LinearModelTest(keras_parameterized.TestCase):
         grads_and_vars = zip(grads, model.trainable_variables)
         opt.apply_gradients(grads_and_vars)
 
+  # This test is an example for a regression on categorical inputs, i.e.,
+  # the output is 0.4, 0.6, 0.9 when input is 'alpha', 'beta', 'gamma'
+  # separately.
+  def test_linear_model_with_feature_column(self):
+    with context.eager_mode():
+      vocab_list = ['alpha', 'beta', 'gamma']
+      vocab_val = [0.4, 0.6, 0.9]
+      data = np.random.choice(vocab_list, size=256)
+      y = np.zeros_like(data, dtype=np.float32)
+      for vocab, val in zip(vocab_list, vocab_val):
+        indices = np.where(data == vocab)
+        y[indices] = val + np.random.uniform(
+            low=-0.01, high=0.01, size=indices[0].shape)
+      cat_column = fc.categorical_column_with_vocabulary_list(
+          key='symbol', vocabulary_list=vocab_list)
+      ind_column = fc.indicator_column(cat_column)
+      dense_feature_layer = dense_features_v2.DenseFeatures([ind_column])
+      linear_model = linear.LinearModel(
+          use_bias=False, kernel_initializer='zeros')
+      combined = sequential.Sequential([dense_feature_layer, linear_model])
+      opt = gradient_descent.SGD(learning_rate=0.1)
+      combined.compile(opt, 'mse', [])
+      combined.fit(x={'symbol': data}, y=y, batch_size=32, epochs=10)
+      self.assertAllClose([[0.4], [0.6], [0.9]],
+                          combined.layers[1].dense_layers[0].kernel.numpy(),
+                          atol=0.01)
+
   def test_config(self):
     linear_model = linear.LinearModel(units=3, use_bias=True)
     config = linear_model.get_config()

tensorflow/python/keras/premade/wide_deep_test.py

@@ -21,11 +21,13 @@ from __future__ import print_function
 import numpy as np
 
 from tensorflow.python.eager import context
+from tensorflow.python.feature_column import feature_column_v2 as fc
 from tensorflow.python.keras import keras_parameterized
 from tensorflow.python.keras import testing_utils
 from tensorflow.python.keras.engine import input_layer
 from tensorflow.python.keras.engine import sequential
 from tensorflow.python.keras.engine import training
+from tensorflow.python.keras.feature_column import dense_features_v2
 from tensorflow.python.keras.layers import core
 from tensorflow.python.keras.optimizer_v2 import gradient_descent
 from tensorflow.python.keras.premade import linear
@@ -186,6 +188,67 @@ class WideDeepModelTest(keras_parameterized.TestCase):
         run_eagerly=testing_utils.should_run_eagerly())
     wide_deep_model.fit(inputs, output, epochs=50)
 
+  # This test is an example for cases where linear and dnn model accepts
+  # same raw input and same transformed inputs, i.e., the raw input is
+  # categorical, and both linear and dnn model accept one hot encoding.
+  def test_wide_deep_model_with_single_feature_column(self):
+    vocab_list = ['alpha', 'beta', 'gamma']
+    vocab_val = [0.4, 0.6, 0.9]
+    data = np.random.choice(vocab_list, size=256)
+    y = np.zeros_like(data, dtype=np.float32)
+    for vocab, val in zip(vocab_list, vocab_val):
+      indices = np.where(data == vocab)
+      y[indices] = val + np.random.uniform(
+          low=-0.01, high=0.01, size=indices[0].shape)
+    cat_column = fc.categorical_column_with_vocabulary_list(
+        key='symbol', vocabulary_list=vocab_list)
+    ind_column = fc.indicator_column(cat_column)
+    dense_feature_layer = dense_features_v2.DenseFeatures([ind_column])
+    linear_model = linear.LinearModel(
+        use_bias=False, kernel_initializer='zeros')
+    dnn_model = sequential.Sequential([core.Dense(units=1)])
+    wide_deep_model = wide_deep.WideDeepModel(linear_model, dnn_model)
+    combined = sequential.Sequential([dense_feature_layer, wide_deep_model])
+    opt = gradient_descent.SGD(learning_rate=0.1)
+    combined.compile(
+        opt,
+        'mse', [],
+        run_eagerly=testing_utils.should_run_eagerly())
+    combined.fit(x={'symbol': data}, y=y, batch_size=32, epochs=10)
+
+  # This test is an example for cases where linear and dnn model accepts
+  # same raw input but different transformed inputs, i.e,. the raw input is
+  # categorical, and linear model accepts one hot encoding, while dnn model
+  # accepts embedding encoding.
+  def test_wide_deep_model_with_two_feature_columns(self):
+    vocab_list = ['alpha', 'beta', 'gamma']
+    vocab_val = [0.4, 0.6, 0.9]
+    data = np.random.choice(vocab_list, size=256)
+    y = np.zeros_like(data, dtype=np.float32)
+    for vocab, val in zip(vocab_list, vocab_val):
+      indices = np.where(data == vocab)
+      y[indices] = val + np.random.uniform(
+          low=-0.01, high=0.01, size=indices[0].shape)
+    cat_column = fc.categorical_column_with_vocabulary_list(
+        key='symbol', vocabulary_list=vocab_list)
+    ind_column = fc.indicator_column(cat_column)
+    emb_column = fc.embedding_column(cat_column, dimension=5)
+    linear_feature_layer = dense_features_v2.DenseFeatures([ind_column])
+    linear_model = linear.LinearModel(
+        use_bias=False, kernel_initializer='zeros')
+    combined_linear = sequential.Sequential(
+        [linear_feature_layer, linear_model])
+    dnn_model = sequential.Sequential([core.Dense(units=1)])
+    dnn_feature_layer = dense_features_v2.DenseFeatures([emb_column])
+    combined_dnn = sequential.Sequential([dnn_feature_layer, dnn_model])
+    wide_deep_model = wide_deep.WideDeepModel(combined_linear, combined_dnn)
+    opt = gradient_descent.SGD(learning_rate=0.1)
+    wide_deep_model.compile(
+        opt,
+        'mse', [],
+        run_eagerly=testing_utils.should_run_eagerly())
+    wide_deep_model.fit(x={'symbol': data}, y=y, batch_size=32, epochs=10)
+
   def test_config(self):
     linear_model = linear.LinearModel(units=1)
     dnn_model = sequential.Sequential([core.Dense(units=1, input_dim=3)])

tensorflow/python/keras/saving/BUILD

@@ -109,6 +109,7 @@ tf_py_test(
     python_version = "PY3",
     deps = [
         "//tensorflow/python:client_testlib",
+        "//tensorflow/python/feature_column:feature_column_v2",
         "//tensorflow/python/keras",
         "//tensorflow/python/keras:combinations",
         "//third_party/py/numpy",

tensorflow/python/keras/saving/save_test.py

@@ -25,14 +25,19 @@ from absl.testing import parameterized
 import numpy as np
 
 from tensorflow.python import keras
+from tensorflow.python.eager import context
+from tensorflow.python.feature_column import feature_column_lib
+from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.framework import test_util
 from tensorflow.python.keras import combinations
 from tensorflow.python.keras import losses
 from tensorflow.python.keras import testing_utils
 from tensorflow.python.keras.engine import sequential
 from tensorflow.python.keras.layers import core
+from tensorflow.python.keras.saving import model_config
 from tensorflow.python.keras.saving import save
 from tensorflow.python.keras.utils import generic_utils
+from tensorflow.python.ops import lookup_ops
 from tensorflow.python.platform import test
 from tensorflow.python.saved_model import loader_impl
 
@@ -100,6 +105,101 @@ class TestSaveModel(test.TestCase, parameterized.TestCase):
       save.save_model(self.model, path, save_format='tf')
       save.load_model(path)
 
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def test_saving_with_dense_features(self):
+    cols = [
+        feature_column_lib.numeric_column('a'),
+        feature_column_lib.indicator_column(
+            feature_column_lib.categorical_column_with_vocabulary_list(
+                'b', ['one', 'two']))
+    ]
+    input_layers = {
+        'a': keras.layers.Input(shape=(1,), name='a'),
+        'b': keras.layers.Input(shape=(1,), name='b', dtype='string')
+    }
+
+    fc_layer = feature_column_lib.DenseFeatures(cols)(input_layers)
+    output = keras.layers.Dense(10)(fc_layer)
+
+    model = keras.models.Model(input_layers, output)
+
+    model.compile(
+        loss=keras.losses.MSE,
+        optimizer='rmsprop',
+        metrics=[keras.metrics.categorical_accuracy])
+
+    config = model.to_json()
+    loaded_model = model_config.model_from_json(config)
+
+    inputs_a = np.arange(10).reshape(10, 1)
+    inputs_b = np.arange(10).reshape(10, 1).astype('str')
+
+    with self.cached_session():
+      # Initialize tables for V1 lookup.
+      if not context.executing_eagerly():
+        self.evaluate(lookup_ops.tables_initializer())
+
+      self.assertLen(loaded_model.predict({'a': inputs_a, 'b': inputs_b}), 10)
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def test_saving_with_sequence_features(self):
+    cols = [
+        feature_column_lib.sequence_numeric_column('a'),
+        feature_column_lib.indicator_column(
+            feature_column_lib.sequence_categorical_column_with_vocabulary_list(
+                'b', ['one', 'two']))
+    ]
+    input_layers = {
+        'a':
+            keras.layers.Input(shape=(None, 1), sparse=True, name='a'),
+        'b':
+            keras.layers.Input(
+                shape=(None, 1), sparse=True, name='b', dtype='string')
+    }
+
+    fc_layer, _ = feature_column_lib.SequenceFeatures(cols)(input_layers)
+    # TODO(tibell): Figure out the right dtype and apply masking.
+    # sequence_length_mask = array_ops.sequence_mask(sequence_length)
+    # x = keras.layers.GRU(32)(fc_layer, mask=sequence_length_mask)
+    x = keras.layers.GRU(32)(fc_layer)
+    output = keras.layers.Dense(10)(x)
+
+    model = keras.models.Model(input_layers, output)
+
+    model.compile(
+        loss=keras.losses.MSE,
+        optimizer='rmsprop',
+        metrics=[keras.metrics.categorical_accuracy])
+
+    config = model.to_json()
+    loaded_model = model_config.model_from_json(config)
+
+    batch_size = 10
+    timesteps = 1
+
+    values_a = np.arange(10, dtype=np.float32)
+    indices_a = np.zeros((10, 3), dtype=np.int64)
+    indices_a[:, 0] = np.arange(10)
+    inputs_a = sparse_tensor.SparseTensor(indices_a, values_a,
+                                          (batch_size, timesteps, 1))
+
+    values_b = np.zeros(10, dtype=np.str)
+    indices_b = np.zeros((10, 3), dtype=np.int64)
+    indices_b[:, 0] = np.arange(10)
+    inputs_b = sparse_tensor.SparseTensor(indices_b, values_b,
+                                          (batch_size, timesteps, 1))
+
+    with self.cached_session():
+      # Initialize tables for V1 lookup.
+      if not context.executing_eagerly():
+        self.evaluate(lookup_ops.tables_initializer())
+
+      self.assertLen(
+          loaded_model.predict({
+              'a': inputs_a,
+              'b': inputs_b
+          }, steps=1), batch_size)
+
   @combinations.generate(combinations.combine(mode=['graph', 'eager']))
   def test_saving_h5_for_rnn_layers(self):
     # See https://github.com/tensorflow/tensorflow/issues/35731 for details.