allow grain to read tfds config

src/maxtext/configs/types.py

@@ -4101,17 +4101,34 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         raise ValueError("Only supports <= 1 for now, more workers results in duplicated data")
     elif self.dataset_type == DatasetType.GRAIN:
       use_hf_parquet = self.hf_path and self.grain_file_type == "parquet"
+      use_tfds_tfrecord_train = (
+          self.grain_file_type == "tfrecord" and self.dataset_path and self.dataset_name and self.train_split
+      )
+      use_tfds_tfrecord_eval = (
+          self.grain_file_type == "tfrecord" and self.dataset_path and self.eval_dataset_name and self.eval_split
+      )
 
-      if not self.grain_train_files and not self.grain_train_mixture_config_path and not use_hf_parquet:
+      if (
+          not self.grain_train_files
+          and not self.grain_train_mixture_config_path
+          and not use_hf_parquet
+          and not use_tfds_tfrecord_train
+      ):
         raise ValueError(
             "When dataset_type=grain, set grain_train_files, "
-            "grain_train_mixture_config_path, or use hf_path with grain_file_type=parquet."
+            "grain_train_mixture_config_path, use hf_path with grain_file_type=parquet, or use dataset_path, "
+            "dataset_name, and train_split with grain_file_type=tfrecord."
+        )
+      if self.eval_interval > 0 and not self.grain_eval_files and not use_hf_parquet and not use_tfds_tfrecord_eval:
+        raise ValueError(
+            "Please specify grain_eval_files, use hf_path with grain_file_type=parquet, or use dataset_path, "
+            "eval_dataset_name, and eval_split with grain_file_type=tfrecord; otherwise set eval_interval to <=0."
         )
-      if self.eval_interval > 0 and not self.grain_eval_files and not use_hf_parquet:
-        raise ValueError("Please specify grain_eval_files (or hf_path with parquet) or set eval_interval to <=0.")
     elif self.dataset_type == DatasetType.TFDS:
       logger.warning(
-          "tfds pipeline is deprecated. Use dataset_type=grain, grain_file_type=tfrecord, and provide grain_train_files."
+          "tfds pipeline is deprecated. Use dataset_type=grain and grain_file_type=tfrecord. You can keep the same "
+          "dataset_path, dataset_name, train_split, eval_dataset_name, and eval_split settings to automatically construct "
+          "the file paths. Alternatively, provide grain_train_files or grain_eval_files for custom file paths."
       )
       if self.use_dpo:
         raise ValueError(

src/maxtext/input_pipeline/grain_data_processing.py

@@ -66,6 +66,14 @@ def construct_hf_dataset_path(hf_path: str, hf_train_files: str | None = None, s
   return full_path
 
 
+def construct_tfds_tfrecord_path(dataset_path: str, dataset_name: str, split: str) -> str:
+  """Constructs a glob for TFRecords in the standard TFDS prepared-data layout."""
+  dataset_dir = dataset_name.strip().strip("/").replace(":", "/")
+  path = f"{dataset_path.strip().rstrip('/')}/{dataset_dir}/*-{split}.tfrecord-*"
+  max_logging.log(f"Automatically constructed Grain TFRecord path from TFDS configuration: {path}")
+  return path
+
+
 def find_data_files(data_file_pattern, hf_access_token=None):
   """Find data files matching the pattern."""
   if data_file_pattern.startswith("gs://"):
@@ -477,8 +485,22 @@ def make_grain_train_iterator(
   pipeline_fn = _get_pipeline_fn(config)
 
   grain_train_files = config.grain_train_files
-  if not grain_train_files and not config.grain_train_mixture_config_path and config.hf_path:
+  if (
+      not grain_train_files
+      and not config.grain_train_mixture_config_path
+      and config.grain_file_type == "parquet"
+      and config.hf_path
+  ):
     grain_train_files = construct_hf_dataset_path(config.hf_path, split="train")
+  elif (
+      not grain_train_files
+      and not config.grain_train_mixture_config_path
+      and config.grain_file_type == "tfrecord"
+      and config.dataset_path
+      and config.dataset_name
+      and config.train_split
+  ):
+    grain_train_files = construct_tfds_tfrecord_path(config.dataset_path, config.dataset_name, config.train_split)
 
   get_ds_fn = functools.partial(
       get_datasets,
@@ -583,9 +605,17 @@ def make_grain_eval_iterator(
   pipeline_fn = _get_pipeline_fn(config)
 
   grain_eval_files = config.grain_eval_files
-  if not grain_eval_files and getattr(config, "hf_path", None):
+  if not grain_eval_files and config.grain_file_type == "parquet" and getattr(config, "hf_path", None):
     split = getattr(config, "hf_eval_split", None) or "validation"
     grain_eval_files = construct_hf_dataset_path(config.hf_path, split=split)
+  elif (
+      not grain_eval_files
+      and config.grain_file_type == "tfrecord"
+      and config.dataset_path
+      and config.eval_dataset_name
+      and config.eval_split
+  ):
+    grain_eval_files = construct_tfds_tfrecord_path(config.dataset_path, config.eval_dataset_name, config.eval_split)
 
   get_ds_fn = functools.partial(
       get_datasets,

src/maxtext/input_pipeline/tfds_data_processing.py

@@ -24,7 +24,12 @@
   import tensorflow_datasets as tfds
 except ImportError as error:
   raise ImportError(
-      "TensorFlow and tensorflow-datasets are required. Run `pip install tensorflow tensorflow-datasets`"
+      "The deprecated TFDS pipeline requires TensorFlow and tensorflow-datasets.\n\n"
+      "Recommended: migrate to Grain by setting dataset_type=grain and grain_file_type=tfrecord. Existing "
+      "dataset_path, dataset_name, train_split, eval_dataset_name, and eval_split settings will be used to "
+      "automatically construct file paths when grain_train_files or grain_eval_files are not provided. Set the "
+      "grain_train_files or grain_eval_files explicitly for a custom path.\n\n"
+      "To continue using the deprecated TFDS pipeline, run: `pip install tensorflow tensorflow-datasets`."
   ) from error
 
 import jax

tests/unit/grain_data_processing_test.py

@@ -34,6 +34,27 @@
 from tests.utils.test_helpers import get_test_base_output_directory, get_test_config_path, get_test_dataset_path
 
 
+class TestTfdsTfrecordPathFallback:
+  """Tests TFDS prepared-data path construction without reading a dataset."""
+
+  def test_construct_tfds_tfrecord_path(self, mocker):
+    log = mocker.patch.object(grain_data_processing.max_logging, "log")
+    assert (
+        grain_data_processing.construct_tfds_tfrecord_path("  gs://maxtext-dataset/  ", "c4/en:3.0.1", "train")
+        == "gs://maxtext-dataset/c4/en/3.0.1/*-train.tfrecord-*"
+    )
+    log.assert_called_once_with(
+        "Automatically constructed Grain TFRecord path from TFDS configuration: "
+        "gs://maxtext-dataset/c4/en/3.0.1/*-train.tfrecord-*"
+    )
+
+  def test_missing_derived_path_reports_pattern(self, tmp_path):
+    pattern = str(tmp_path / "c4/en/3.0.1/*-train.tfrecord-*")
+
+    with pytest.raises(FileNotFoundError, match=r"No files found matching pattern: .*\*-train\.tfrecord-\*"):
+      grain_data_processing.find_data_files(pattern)
+
+
 class GrainBaseProcessingTest:
   """Base mixin with test_train_ds for all grain data processing tests.
 
@@ -532,6 +553,16 @@ class GrainTFRecordProcessingTest(_GrainTFRecordSetup, GrainDeterminismMixin, Gr
   def setUpClass(cls):
     super().setUpClass()
 
+  def test_config_accepts_tfds_tfrecord_fallback(self):
+    config = self._make_config(
+        grain_train_files="",
+        dataset_path="gs://maxtext-dataset",
+        dataset_name="c4/en:3.0.1",
+        train_split="train",
+        eval_interval=0,
+    )
+    self.assertEqual(config.grain_train_files, "")
+
 
 class GrainTFRecordPreTokenizedProcessingTest(_GrainTFRecordSetup, GrainBaseProcessingTest, unittest.TestCase):
   """Test grain data processing with a pre-tokenized TFRecord dataset (tokenize_train_data=False).