Switch dataset to the new tf.data.Dataset API.
Switch dataset to the new tf.data.Dataset API.

PiperOrigin-RevId: 369679305

deepmind/deepmind-research,deepmind/deepmind-research,deepmind/deepmind-research,deepmind/deepmind-research