@markdaoust here’s what I tried:
def class_func(image, label):
    return label

SUPPORT_BATCH_SIZE = 640

(x_train, y_train), (_, _) = tf.keras.datasets.cifar10.load_data()

sampled_idx = np.random.choice(len(x_train), 4000)
sampled_train, sampled_labels = x_train[sampled_idx], y_train[sampled_idx].squeeze()
sampled_labels = sampled_labels.astype("int32")
support_ds = tf.data.Dataset.from_tensor_slices((sampled_train, sampled_labels))

distribution = Counter(sampled_labels)
counts = np.array(list(distribution.values()))
fractions = counts/counts.sum().astype("float64")

target_distribution = np.array([0.1] * 10).astype("float64")
resampler = tf.data.experimental.rejection_resample(
    class_func, target_dist=target_distribution, initial_dist=fractions)
support_ds = support_ds.apply(resampler).batch(SUPPORT_BATCH_SIZE)

Here’s the root error:
TypeError: Input 'y' of 'Less' Op has type float64 that does not match type float32 of argument 'x'.

Any idea what I might have missed out on? Here’s the Colab 8 if you wanna give it a shot.