Function train_test_split 

pub fn train_test_split(
    data: &Array2<f64>,
    labels: &Array1<f64>,
    test_size: f64,
    shuffle: bool,
) -> (Array2<f64>, Array2<f64>, Array1<f64>, Array1<f64>)

Randomly split data into train and test sets