first_features, first_labels = next(iter(dataset))
row_count = sum(int(labels.shape[0]) for _, labels in dataset)

assert set(first_features) == {"feature_a", "feature_b"}
assert first_labels.numpy().tolist() == [0, 1]
assert row_count == 4

for name, values in first_features.items():
    print(f"{name}={values.numpy()}")
print(f"labels={first_labels.numpy()}")
print(f"rows={row_count}")