import tensorflow as tf

features = tf.reshape(tf.cast(tf.range(40), tf.float32), (20, 2))
labels = tf.cast(tf.range(20) % 2, tf.int32)

full_dataset = tf.data.Dataset.from_tensor_slices((features, labels))

train_split, holdout_split = tf.keras.utils.split_dataset(
    full_dataset,
    left_size=0.7,
    shuffle=True,
    seed=7,
)
validation_split, test_split = tf.keras.utils.split_dataset(
    holdout_split,
    left_size=0.5,
    shuffle=False,
)

train_count = int(train_split.cardinality().numpy())
validation_count = int(validation_split.cardinality().numpy())
test_count = int(test_split.cardinality().numpy())

train_ds = train_split.batch(4)
validation_ds = validation_split.batch(4)
test_ds = test_split.batch(4)

train_features, train_labels = next(iter(train_ds))
validation_features, validation_labels = next(iter(validation_ds))
test_features, test_labels = next(iter(test_ds))

print(f"train_examples={train_count}")
print(f"validation_examples={validation_count}")
print(f"test_examples={test_count}")
print(f"total_examples={train_count + validation_count + test_count}")
print(f"train_batch_shape={tuple(train_features.shape)}")
print(f"train_label_shape={tuple(train_labels.shape)}")
print(f"validation_batch_shape={tuple(validation_features.shape)}")
print(f"validation_label_shape={tuple(validation_labels.shape)}")
print(f"test_batch_shape={tuple(test_features.shape)}")
print(f"test_label_shape={tuple(test_labels.shape)}")