TFDS This shows how to interface with JAX and PyTorch https://www.tensorflow.org/datasets/tfless_tfds