How to do train_test split in stratified manner using ray?

1

What is the equivalent command to sklearn.model_selection.train_test_split that can stratify the samples based on class labels?
Th ray.data.Dataset.train_test_split does have stratify option.

How can I achieve stratified split in ray data?