Построение KNeighborsRegressor для разных значений K

Question

Я хотел бы построить график KNeighborsRegressor для различных значений K.

Я попробовал этот код (я восстановил его из своих заметок):

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2, random_state = 0)

fig, subaxes = plt.subplots(5, 1, figsize=(5,20))
X_predict_input = np.linspace(-3, 3, 500).reshape(-1,1)


for thisaxis, K in zip(subaxes, [1, 5, 10, 13, 20, 55]):  
    knnreg = KNeighborsRegressor(n_neighbors = K).fit(X_train, y_train)
    y_predict_output = knnreg.predict(X_predict_input)
    train_score = knnreg.score(X_train, y_train)
    test_score = knnreg.score(X_test, y_test)
    thisaxis.plot(X_predict_input, y_predict_output)
    thisaxis.plot(X_train, y_train, 'o', alpha=0.9, label='Train')
    thisaxis.plot(X_test, y_test, '^', alpha=0.9, label='Test')
    thisaxis.set_xlabel('Input feature')
    thisaxis.set_ylabel('Target value')
    thisaxis.set_title('KNN Regression (K={})\n\ Train $R^2 = {:.3f}$,  Test $R^2 = {:.3f}$'.format(K, train_score, test_score))
    thisaxis.legend()
    plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=1.0)

Где

X_train.shape, X_test.shape

равны (1213,21) и (304,21) соответственно.

Но я получил следующая ошибка:

ValueError: query data dimension must match training data dimension

Я думаю, проблема в переменной X_predict_input , но я не могу это исправить.