Introduce Tensorflow Dataset

[BillHuang2001]

Description

Introduce Tensorflow Dataset and deprecate the old torchvision dataset.

Motivation:

1. Torchvision requires torch to co-exist with JAX, which is a pain to manage. Tensorflow dataset is mostly standalone (download the dataset require tensorflow, if already downloaded, tensorflow is not required), and tensorflow/jax co-existence is much simpler.
2. Tensorflow dataset has a much larger dataset coverage, not only vision tasks.
3. The default data format is the same as flax/haiku. For example, the default format for image is (N, H, W, C), while torchvision gives (N, C, H, W).
4. Torch's dataloader is really bad.

Checklist

• I have formatted my Python code with black.
• I have good commit messages.
• If adding new algorithms, problems, operators:
  • Added related test cases.
  • Added docstring to explain important parameters.
  • Added entries in the docs.