Оболочка для test_train_split для создания обучающих, проверочных и тестовых разделений для любого количества входных массивов

Сигнатура функции:

train_test_split(*arrays, **options)

это означает, что он принимает любое количество позиционных массивов и любое количество параметров ключевых слов. Чтобы вернуть train, val, test по вашему желанию, нужно действовать следующим образом:

from sklearn.model_selection import train_test_split
df = pd.DataFrame({"x": np.random.randn(1000),"y": np.random.randn(1000)})

def train_validate_test_split(dataframe, **options):
   train, test = train_test_split(dataframe, **options)
   train, val = train_test_split(train, **options)
   return train, val, test

a,b,c = train_validate_test_split(df, train_size=.25)

РЕДАКТИРОВАТЬ

Чтобы принять 1 или 2 входа, используйте:

def train_val_test_split(*arrays,**options):

    if len(arrays) == 1:
        X_train, X_test = train_test_split(*arrays,**options)
        X_train, X_val = train_test_split(X_train,**options)
        print("Unpack to X_train, X_val, X_test")
        return X_train, X_val, X_test

    if len(arrays) == 2:
        X_train, X_test, y_train, y_test = train_test_split(*arrays,**options)
        X_train, X_val, y_train, y_val = train_test_split(X_train,y_train,**options)
        print("Unpack to X_train, X_val, X_test, y_train, y_val, y_test")
        return X_train, X_val, X_test, y_train, y_val, y_test

    else:
        raise ValueError("Only implemented for 1 or 2 arrays. "
                          f"You provided {len(arrays)} arrays")

или для любого количества входных массивов:

y = np.random.randn(1000)
def train_val_test_split(*arrays,**options):
    '''
    inputs:
        arrays - any number of array to split,
    outputs:
        sequence 
        arr1_train, arr2_train, ... , arr1_val , arr2_val, ..., arr1_test, arr2_test, ...
    '''
    *out, = train_test_split(*arrays,**options)
    train = out[0::2] #x1_train, x2_train, ...
    test  = out[1::2] #x1_test, x2_test, ...
    *train_val, = train_test_split(*train,**options)
    train = train_val[0::2]
    val   = train_val[1::2]
    print(f"Unpack to {len(arrays)*3} tuples: train,...,val,..., test...")
    return tuple(split for tuple_ in zip(train,val,test) for split in tuple_)

x = train_val_test_split(y,y,y)

for item in x:
    print(item.shape, end = ", ")

Unpack to 9 tuples: train,...,val,..., test...
(562,), (188,), (250,), (562,), (188,), (250,), (562,), (188,), (250,),