使用tensorflow将数据集拆分为训练和测试

我想将我的完整数据集（每个原始数据都有多个特性）拆分为训练集和测试集。除了使用scikit learn的列车测试拆分，还有其他合适的方法拆分我的数据吗？以及我需要洗牌我的数据时，分裂。 （如果建议的方法是基于tensorflow的，那就太好了。）

试试下面的代码：

import tensorflow as tf
input = tf.random.uniform([100, 5], 0, 10, dtype=tf.int32)
input = tf.random.shuffle(input)
train_ds = input[:90]
test_ds = input[-10:]