有没有办法用tf.data.Dataset做一个分层的随机采样器？

import tensorflow as tf

# 假设有三个层级的数据集，分别为dataset1、dataset2、dataset3

# 转换为Dataset对象
dataset1 = tf.data.Dataset.from_tensor_slices(data1)
dataset2 = tf.data.Dataset.from_tensor_slices(data2)
dataset3 = tf.data.Dataset.from_tensor_slices(data3)

# 对每个层级的数据集进行随机化操作
dataset1 = dataset1.shuffle(buffer_size=len(data1))
dataset2 = dataset2.shuffle(buffer_size=len(data2))
dataset3 = dataset3.shuffle(buffer_size=len(data3))

# 重复数据集
dataset1 = dataset1.repeat()
dataset2 = dataset2.repeat()
dataset3 = dataset3.repeat()

# 划分批次
dataset1 = dataset1.batch(batch_size)
dataset2 = dataset2.batch(batch_size)
dataset3 = dataset3.batch(batch_size)

# 提前加载数据
dataset1 = dataset1.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
dataset2 = dataset2.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
dataset3 = dataset3.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)