Merge pull request #15473 from diggerk:index_lookup_sparse_config

PiperOrigin-RevId: 401828338

Author: tensorflower-gardener

keras/layers/preprocessing/index_lookup.py

@@ -363,6 +363,7 @@ def get_config(self):
         "oov_token": self.oov_token,
         "mask_token": self.mask_token,
         "output_mode": self.output_mode,
+        "sparse": self.sparse,
         "pad_to_max_tokens": self.pad_to_max_tokens,
         "vocabulary": utils.listify_tensors(self.input_vocabulary),
         "vocabulary_dtype": self.vocabulary_dtype,

keras/layers/preprocessing/index_lookup_test.py

@@ -2365,6 +2365,30 @@ def test_static_table_config_weight_data_transfer_succeeds(self):
     new_output_dataset = model.predict(input_array)
     self.assertAllEqual(new_output_dataset, expected_output)
 
+  def test_sparse_output_across_saving(self):
+    vocab_data = ["earth", "wind", "and", "fire"]
+    input_array = np.array([["earth", "wind", "and", "fire"],
+                            ["fire", "and", "earth", "michigan"]])
+
+    expected_output = [[0., 1., 1., 1., 1.], [1., 1., 0., 1., 1.]]
+
+    layer_cls = index_lookup.IndexLookup
+    layer = layer_cls(
+        max_tokens=None,
+        num_oov_indices=1,
+        mask_token="",
+        oov_token="[OOV]",
+        vocabulary_dtype=tf.string,
+        vocabulary=vocab_data,
+        output_mode="multi_hot",
+        sparse=True)
+    config = layer.get_config()
+    layer = layer_cls.from_config(config)
+
+    output = layer(input_array)
+    self.assertIsInstance(output, tf.SparseTensor)
+    self.assertAllEqual(tf.sparse.to_dense(output), expected_output)
+
 
 class EagerExecutionDisabled(keras_parameterized.TestCase,
                              preprocessing_test_utils.PreprocessingLayerTest):