Apply sharding when creating tensors data in data loader (#5015)
* Revert "Revert "Apply sharding when creating tensors data in data loader (#4995)" (#5012)"
This reverts commit 6e9b663c309a0db079d7bc354361d9e6e852b0b3.
* Change test tensor shape to support TPU v3
