tf-models-nightly 2.17.0.dev20240617__py2.py3-none-any.whl → 2.20.0.dev20251205__py2.py3-none-any.whl

official/projects/yolo/tasks/task_utils.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yolo/tasks/yolo.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yolo/tasks/yolov7.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yolo/train.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/configs/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/configs/yt8m.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/configs/yt8m_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/backbones/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/backbones/dbof.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/backbones/dbof_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/heads/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/heads/logistic.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/heads/moe.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/nn_layers.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/nn_layers_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/yt8m_model.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/yt8m_model_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/yt8m_model_utils.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/modeling/yt8m_model_utils_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/tasks/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/tasks/yt8m_task.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/train.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/projects/yt8m/train_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/constants.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/create_ncf_data.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -17,7 +17,6 @@
 import json
 
 # pylint: disable=g-bad-import-order
-# Import libraries
 from absl import app
 from absl import flags
 import tensorflow as tf, tf_keras

official/recommendation/data_pipeline.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/data_preprocessing.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/data_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -190,7 +190,7 @@ class BaseTest(tf.test.TestCase):
         train_examples[l].add((u_raw, i_raw))
         counts[(u_raw, i_raw)] += 1
 
-    self.assertRegexpMatches(md5.hexdigest(), END_TO_END_TRAIN_MD5)
+    self.assertRegex(md5.hexdigest(), END_TO_END_TRAIN_MD5)
 
     num_positives_seen = len(train_examples[True])
     self.assertEqual(producer._train_pos_users.shape[0], num_positives_seen)
@@ -254,7 +254,7 @@ class BaseTest(tf.test.TestCase):
           # from the negatives.
           assert (u_raw, i_raw) not in self.seen_pairs
 
-    self.assertRegexpMatches(md5.hexdigest(), END_TO_END_EVAL_MD5)
+    self.assertRegex(md5.hexdigest(), END_TO_END_EVAL_MD5)
 
   def _test_fresh_randomness(self, constructor_type):
     train_epochs = 5
@@ -300,7 +300,7 @@ class BaseTest(tf.test.TestCase):
         else:
           negative_counts[(u, i)] += 1
 
-    self.assertRegexpMatches(md5.hexdigest(), FRESH_RANDOMNESS_MD5)
+    self.assertRegex(md5.hexdigest(), FRESH_RANDOMNESS_MD5)
 
     # The positive examples should appear exactly once each epoch
     self.assertAllEqual(

official/recommendation/movielens.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -27,7 +27,6 @@ import tempfile
 import zipfile
 
 # pylint: disable=g-bad-import-order
-# Import libraries
 import numpy as np
 import pandas as pd
 import six

official/recommendation/ncf_common.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/ncf_input_pipeline.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/ncf_keras_main.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/ncf_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/neumf_model.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/popen_helper.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/ranking/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/ranking/common.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/ranking/configs/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/ranking/configs/config.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -74,6 +74,7 @@ class DataConfig(hyperparams.Config):
   cycle_length: int = 10
   sharding: bool = True
   num_shards_per_host: int = 8
+  use_cached_data: bool = False
 
 
 @dataclasses.dataclass
@@ -112,6 +113,11 @@ class ModelConfig(hyperparams.Config):
       module
     dcn_use_bias: Flag to determine whether to use bias for the dcn interaction
       module
+    use_partial_tpu_embedding: Flag to determine whether to use partial tpu
+      embedding layer or not.
+    max_ids_per_chip_per_sample: Maximum number of ids per chip per sample.
+    max_ids_per_table: Maximum number of ids per table.
+    max_unique_ids_per_table: Maximum number of unique ids per table.
   """
   num_dense_features: int = 13
   vocab_sizes: List[int] = dataclasses.field(default_factory=list)
@@ -128,6 +134,12 @@ class ModelConfig(hyperparams.Config):
   dcn_kernel_initializer: str = 'truncated_normal'
   dcn_bias_initializer: str = 'zeros'
   dcn_use_bias: bool = True
+  use_partial_tpu_embedding: bool = True
+  max_ids_per_chip_per_sample: int | None = None
+  max_ids_per_table: Union[int, List[int]] | None = None
+  max_unique_ids_per_table: Union[int, List[int]] | None = None
+  allow_id_dropping: bool = False
+  initialize_tables_on_host: bool = False
 
 
 @dataclasses.dataclass
@@ -424,6 +436,7 @@ def dlrm_dcn_v2_criteo_tb_config() -> Config:
               dcn_use_bias=True,
               concat_dense=False,
               use_multi_hot=True,
+              use_partial_tpu_embedding=False,
               multi_hot_sizes=multi_hot_sizes,
               ),
           loss=Loss(label_smoothing=0.0),

official/recommendation/ranking/configs/config_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/ranking/data/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/recommendation/ranking/data/data_pipeline.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -38,12 +38,14 @@ class CriteoTsvReader:
                params: config.DataConfig,
                num_dense_features: int,
                vocab_sizes: List[int],
-               use_synthetic_data: bool = False):
+               use_synthetic_data: bool = False,
+               use_cached_data: bool = False):
     self._file_pattern = file_pattern
     self._params = params
     self._num_dense_features = num_dense_features
     self._vocab_sizes = vocab_sizes
     self._use_synthetic_data = use_synthetic_data
+    self._use_cached_data = use_cached_data
 
   def __call__(self, ctx: tf.distribute.InputContext) -> tf.data.Dataset:
     params = self._params
@@ -117,6 +119,8 @@ class CriteoTsvReader:
         num_parallel_calls=tf.data.experimental.AUTOTUNE)
 
     dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
+    if self._use_cached_data:
+      dataset = dataset.take(1).cache().repeat()
 
     return dataset
 
@@ -173,6 +177,9 @@ class CriteoTsvReader:
     if params.is_training:
       dataset = dataset.repeat()
 
+    if self._use_cached_data:
+      dataset = dataset.take(1).cache().repeat()
+
     return dataset.batch(batch_size, drop_remainder=True)
 
 

official/recommendation/ranking/data/data_pipeline_multi_hot.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -144,6 +144,8 @@ class CriteoTsvReaderMultiHot:
         num_parallel_calls=tf.data.experimental.AUTOTUNE)
 
     dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
+    if self._params.use_cached_data:
+      dataset = dataset.take(1).cache().repeat()
 
     return dataset
 
@@ -215,12 +217,14 @@ class CriteoTFRecordReader(object):
                params: config.DataConfig,
                num_dense_features: int,
                vocab_sizes: List[int],
-               multi_hot_sizes: List[int],):
+               multi_hot_sizes: List[int],
+               use_cached_data: bool = False):
     self._file_pattern = file_pattern
     self._params = params
     self._num_dense_features = num_dense_features
     self._vocab_sizes = vocab_sizes
     self._multi_hot_sizes = multi_hot_sizes
+    self._use_cached_data = use_cached_data
 
     self.label_features = 'label'
     self.dense_features = ['dense-feature-%d' % x for x in range(1, 14)]
@@ -307,6 +311,8 @@ class CriteoTFRecordReader(object):
         num_parallel_calls=tf.data.experimental.AUTOTUNE,
     )
     dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+    if self._use_cached_data:
+      dataset = dataset.take(1).cache().repeat()
 
     return dataset
 

official/recommendation/ranking/data/data_pipeline_multi_hot_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -23,9 +23,13 @@ from official.recommendation.ranking.data import data_pipeline_multi_hot
 
 class DataPipelineTest(parameterized.TestCase, tf.test.TestCase):
 
-  @parameterized.named_parameters(('Train', True),
-                                  ('Eval', False))
-  def testSyntheticDataPipeline(self, is_training):
+  @parameterized.named_parameters(
+      ('TrainCached', True, True),
+      ('EvalNotCached', False, False),
+      ('TrainNotCached', True, False),
+      ('EvalCached', False, True),
+  )
+  def testSyntheticDataPipeline(self, is_training, use_cached_data):
     task = config.Task(
         model=config.ModelConfig(
             embedding_dim=4,
@@ -39,8 +43,10 @@ class DataPipelineTest(parameterized.TestCase, tf.test.TestCase):
             dcn_low_rank_dim=64,
             bottom_mlp=[64, 32, 4],
             top_mlp=[64, 32, 1]),
-        train_data=config.DataConfig(global_batch_size=16),
-        validation_data=config.DataConfig(global_batch_size=16),
+        train_data=config.DataConfig(global_batch_size=16,
+                                     use_cached_data=use_cached_data),
+        validation_data=config.DataConfig(global_batch_size=16,
+                                          use_cached_data=use_cached_data),
         use_synthetic_data=True)
 
     num_dense_features = task.model.num_dense_features

official/recommendation/ranking/data/data_pipeline_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -23,19 +23,29 @@ from official.recommendation.ranking.data import data_pipeline
 
 class DataPipelineTest(parameterized.TestCase, tf.test.TestCase):
 
-  @parameterized.named_parameters(('Train', True),
-                                  ('Eval', False))
-  def testSyntheticDataPipeline(self, is_training):
+  @parameterized.named_parameters(
+      ('TrainCached', True, True),
+      ('EvalNotCached', False, False),
+      ('TrainNotCached', True, False),
+      ('EvalCached', False, True),
+  )
+  def testSyntheticDataPipeline(self, is_training, use_cached_data):
     task = config.Task(
         model=config.ModelConfig(
             embedding_dim=4,
             num_dense_features=8,
             vocab_sizes=[40, 12, 11, 13, 2, 5],
             bottom_mlp=[64, 32, 4],
-            top_mlp=[64, 32, 1]),
-        train_data=config.DataConfig(global_batch_size=16),
-        validation_data=config.DataConfig(global_batch_size=16),
-        use_synthetic_data=True)
+            top_mlp=[64, 32, 1],
+        ),
+        train_data=config.DataConfig(
+            global_batch_size=16, use_cached_data=use_cached_data
+        ),
+        validation_data=config.DataConfig(
+            global_batch_size=16, use_cached_data=use_cached_data
+        ),
+        use_synthetic_data=True,
+    )
 
     num_dense_features = task.model.num_dense_features
     num_sparse_features = len(task.model.vocab_sizes)