予想ラベルのつけ方がわからない

Question

mnistを読み込ませて訓練データとテストデータに分け、test_Xに対して予想ラベルpred_yを作り, homework関数の戻り値とさせたいです。
今、

def homework(train_X, train_y, test_X):
    # WRITE ME
    mnist = fetch_mldata('MNIST original')
    mnist_X,mnist_y = shuffle(mnist.data,mnist.target,random_state= 42)
    mnist_X = mnist_X/255.0
    train_X,test_X,train_y,test_y =  train_test_split(mnist_X,mnist_y,test_size=0.2,random_state=42)
    train_X,valid_X,train_y,valid_y = train_test_split(train_X,train_y,test_size=0.2,random_state=42)
###########何か追加
    return pred_y

とコードを書いているのですが
test_Xに対してラベルを付ける方法がわかりません。
どう書けばよいのでしょうか？
ちなみに、実行コードは

from sklearn.utils import shuffle
from sklearn.metrics import f1_score
from sklearn.datasets import fetch_mldata
from sklearn.model_selection import train_test_split

import numpy as np

def load_mnist():
    mnist = fetch_mldata('MNIST original')
    mnist_X, mnist_y = shuffle(mnist.data.astype('float32'),
                               mnist.target.astype('int32'), random_state=42)

    mnist_X = mnist_X / 255.0

    return train_test_split(mnist_X, mnist_y,
                test_size=0.2,
                random_state=42)

def validate_homework():
    train_X, test_X, train_y, test_y = load_mnist()

    # validate for small dataset
    train_X_mini = train_X[:1000]
    train_y_mini = train_y[:1000]
    test_X_mini = test_X[:1000]
    test_y_mini = test_y[:1000]

    pred_y = homework(train_X_mini, train_y_mini, test_X_mini)
    print(f1_score(test_y_mini, pred_y, average='macro'))

def score_homework():
    train_X, test_X, train_y, test_y = load_mnist()
    pred_y = homework(train_X, train_y, test_X)
    print(f1_score(test_y, pred_y, average='macro'))

validate_homework()

です。