import numpy as np
from sklearn.base import clone
from sklearn.datasets import load_boston, load_iris
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.model_selection import KFold, StratifiedKFold
from sklearn.model_selection import cross_val_score as skcross_val_score

def cross_val_score(estimator, X, y):
    if estimator._estimator_type == "regressor":
        cv = KFold()
    else:  # estimator._estimator_type == "classifier"
        cv = StratifiedKFold()
    scores = []
    for train, test in cv.split(X, y):
        est = clone(estimator)
        est.fit(X[train], y[train])
        scores.append(est.score(X[test], y[test]))
    return np.array(scores)

# regression
X, y = load_boston(return_X_y=True)
clf = RandomForestRegressor(random_state=0)
ans1 = cross_val_score(clf, X, y)
ans2 = skcross_val_score(clf, X, y)
assert np.allclose(ans1, ans2)

# classification
X, y = load_iris(return_X_y=True)
clf = RandomForestClassifier(random_state=0)
ans1 = cross_val_score(clf, X, y)
ans2 = skcross_val_score(clf, X, y)
assert np.allclose(ans1, ans2)