from pathlib import Path
import shutil

import tensorflow as tf


log_dir = Path("logs/profile-training")
if log_dir.exists():
    shutil.rmtree(log_dir)

features = tf.random.stateless_uniform((512, 32), seed=(7, 11))
labels = tf.cast(tf.reduce_sum(features, axis=1, keepdims=True) > 16.0, tf.float32)
dataset = tf.data.Dataset.from_tensor_slices((features, labels))
dataset = dataset.batch(32).repeat().prefetch(tf.data.AUTOTUNE)