Function train_test_split 

pub fn train_test_split(
    x: &Tensor,
    y: &Tensor,
    test_size: f32,
    shuffle: bool,
) -> (Tensor, Tensor, Tensor, Tensor)

Train-test split utility