试试看：

result = ds.map(lambda x: {
    'image': x['image'],
    'label': x['label']
})

result.element_spec
>>> <MapDataset element_spec={'image': TensorSpec(shape=(32, 32, 3), dtype=tf.uint8, name=None), 'label': TensorSpec(shape=(), dtype=tf.int64, name=None)}>

如果使用tfds.as_numpy，则可以获取字典形式的数据集，并轻松删除列：

import tensorflow_datasets as tfds

ds_builder = tfds.builder('cifar10')
ds_builder.download_and_prepare()
data = tfds.as_numpy(ds_builder.as_dataset(split='train', batch_size=-1))

print(type(data))
# <class 'dict'>

print(data.keys())
# dict_keys(['id', 'image', 'label'])

del data['id']
print(data.keys())
# dict_keys(['image', 'label'])

这也是将数据集导入JAX所需的表单。