Precision @ K как пользовательская метрика в Керасе

Question

Я включаю Precision @ K в качестве пользовательской метрики в Keras. Согласно документации

import keras.backend as K

def mean_pred(y_true, y_pred):
    return K.mean(y_pred)

model.compile(optimizer='rmsprop',
              loss='binary_crossentropy',
              metrics=['accuracy', mean_pred])

Мне просто нужно вычислить функцию, используя бэкэнд Keras, и передать ее на этапе компиляции.

С Numpy только точность @ k может быть вычислена как следующий пример:

def precisionatk(y_true,y_pred,k)
    precision_average = []
    idx =  (-y_pred).argsort(axis=-1)[:,:k]
    for i in range(idx.shape[0]):
        precision_sample = 0
        for j in idx[i,:]:
            if y_true[i,j] == 1:
                precision_sample += 1
        precision_sample = precision_sample / k
        precision_average.append(precision_sample)
    return np.mean(precision_average)


y_true = np.array([[0,0,1,0],[1,0,1,0]])
y_pred = np.array([[0.1,0.4,0.8,0.2],[0.3,0.2,0.5,0.1]])

print(precisionatk(y_true,y_pred,2))

0.75

Итак, как мне перевести это на бэкэнд Keras?

РЕДАКТИРОВАТЬ: Я работаю с проблемой с несколькими метками, и y_true - это всегда массив с единицами или нулями, и y_prediction - вероятность каждого класса.