Как упростить вложенный цикл через Numpy (Наивный Байес)?

Question

Я бы хотел упростить вложенный цикл в одну или две строки с помощью Python numpy. Это последняя часть модели Наивного Байеса в MNIST.

Я понял некоторые функции широковещательной трансляции, но вот сложный цикл с 3 параметрами.

def Bayes_pret(test_x, test_y, prioriP, posteriorP, image_size, num_class):
    pred = np.empty(test_x.shape[0])
    for i in range(test_x.shape[0]): #10000
        prob = np.empty(num_class)
        for j in range(num_class): #10
            temp = sum([(1-test_x[i][x])*math.log(1-posteriorP[j][x]) + test_x[i][x]*math.log(posteriorP[j][x]) for x in range(image_size**2)])
            # where image_size**2 = 784  
            prob[j] = np.array(math.log(prioriP[j]) + temp)
        pred[i] = np.argmax(prob)
        # print('label=',test_y[i],'and pred=',pred[i])
    return pred, (pred == test_y).sum()/ test_y.shape[0]

Это потратило ОГРОМНОЕ время,Я попытался упростить трюк if-else здесь.

temp = sum([math.log(1-posteriorP[j][x]) if test_x[i][x] == 0 else math.log(posteriorP[j][x]) for x in range(image_size**2)])

Но это все еще далеко от совершенства. Я считаю, что не должно быть петли.