[FEATURE] Distributed Data Loader

[merrymercy]

Task 1:

Implement a data loader for the resnet benchmark

• Parax's resnet benchmark
  https://github.com/parax-project/parax/blob/master/benchmark/parax/benchmark_2d_one_case_wresnet.py
• Flax's resnet example
  https://github.com/google/flax/blob/main/examples/imagenet/train.py

class Dataloader:
    def __init__(self,
                 file_path: Union[str, tensorflow.dataloader],
                 physical_mesh: parax.device_mesh.PhysicalMesh,
                 sharding_specs: jax.interpreter.pxla.ShardingSpec,
                 pre_process_func: Callable):
       pass

   def __iter__(self) -> Iteratable[Dict[str -> DistributedArray]]:
       yield batch

Task 2:

Implement a data loader for the BERT benchmark
parax's benchmark
https://github.com/parax-project/parax/blob/master/benchmark/parax/benchmark_2d_one_case_bert.py
huggingface's script
https://github.com/huggingface/transformers/blob/master/examples/flax/language-modeling/run_mlm_flax.py

[merrymercy]

closed by #630