updated to latest code changes in grain

Author: ajkv-google

src/maxtext/input_pipeline/data_processing_utils.py

@@ -26,7 +26,7 @@
 
 def parse_and_keep_features(dataset, config, data_columns, tokenize):
   """Parse arrayrecord features or keep specified columns for other formats."""
-  if config.grain_file_type == "arrayrecord":
+  if config.grain_file_type in ("arrayrecord", "tfrecord"):
     dataset = dataset.map(input_pipeline_utils.ParseFeatures(data_columns, tokenize))
     dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(data_columns, tokenize))
   else: