Tensorflow Dataset API split to train/test

Splitting a tf Dataset is surprisingly hard. Best option is

test_dataset = all_dataset.take(1000) 
train_dataset = all_dataset.skip(1000)

Along with stuff like this if you know the dataset size:

train_size = int(0.7 * DATASET_SIZE)
val_size = int(0.15 * DATASET_SIZE)
test_size = int(0.15 * DATASET_SIZE)

(Though I don’t see any fuction to get a dataset size from the Dataset).

Both from Split a dataset created by Tensorflow dataset API in to Train and Test? - Stack Overflow