linter update

Author: NuojCheng

MaxText/input_pipeline/_syn_data_processing.py

@@ -114,4 +114,4 @@ def get_place_holder_synthetic_data(config: pyconfig.HyperParameters):
     dataset = tf.data.Dataset.zip((output))  # pytype: disable=wrong-arg-types
     dataset = dataset.repeat()
     dataset = dataset.batch(config.global_batch_size_to_load // jax.process_count())
-    return dataset
+    return dataset

MaxText/input_pipeline/input_pipeline_interface.py

@@ -15,19 +15,11 @@
 """
 
 """Input pipeline"""
-from collections.abc import Callable
-from typing import Any
 import functools
 
-import numpy as np
-
-import tensorflow as tf
-
 import jax
-import jax.numpy as jnp
 from jax.sharding import PartitionSpec as P
 
-from MaxText import multihost_dataloading
 from MaxText import pyconfig
 from MaxText import max_logging
 from MaxText.input_pipeline._grain_data_processing import make_grain_train_iterator, make_grain_eval_iterator
@@ -94,7 +86,6 @@ def create_data_iterator(config: pyconfig.HyperParameters, mesh):
     assert len(process_indices_train) == jax.process_count() // config.expansion_factor_real_data
     if config.eval_interval > 0:
       assert len(process_indices_eval) == jax.process_count() // config.expansion_factor_real_data
-  
   # Generate iterator functions according to dataset type
   if config.dataset_type in ["tfds", "grain", "hf", "c4_mlperf"]:
     if config.dataset_type == "c4_mlperf":