import tensorflow as tf

features = {
    "sepal_length": tf.constant([5.1, 4.9, 6.7, 5.6], dtype=tf.float32),
    "sepal_width": tf.constant([3.5, 3.0, 3.1, 2.8], dtype=tf.float32),
}
labels = tf.constant([0, 0, 2, 1], dtype=tf.int32)

dataset = tf.data.Dataset.from_tensor_slices((features, labels))
batched = dataset.batch(2)

first_features, first_label = next(iter(dataset))
batch_features, batch_labels = next(iter(batched))
feature_specs, label_spec = dataset.element_spec

rounded_first_features = {
    name: round(float(value.numpy()), 1) for name, value in first_features.items()
}
rounded_batch_features = {
    name: [round(float(item), 1) for item in value.numpy().tolist()]
    for name, value in batch_features.items()
}
print(f"tensorflow={tf.__version__}")
print(f"dataset_cardinality={int(dataset.cardinality())}")
print(
    "sepal_length_spec="
    f"shape={feature_specs['sepal_length'].shape}, "
    f"dtype={feature_specs['sepal_length'].dtype.name}"
)
print(
    "sepal_width_spec="
    f"shape={feature_specs['sepal_width'].shape}, "
    f"dtype={feature_specs['sepal_width'].dtype.name}"
)
print(f"label_spec=shape={label_spec.shape}, dtype={label_spec.dtype.name}")
print(f"first_example={rounded_first_features}, label={first_label.numpy().item()}")
print(f"first_batch={rounded_batch_features}, labels={batch_labels.numpy().tolist()}")