I am trying to make the following function as fast as possible on large matrices. The matrices can have dimensions in the range of 10K-100K. The matrix values are always between 0 and 1. The function resembles matrix multiplication, but with log operations in the inner loop. These log operations appear to be a bottleneck.
I tried various ways of using Numba and Cython. I also tried writing as much as I could with Numpy. I also experimented with doing fewer memory lookups, but this did not seem to give much advantage. The fastest version is below. High precision is greatly preferred, but if there is a way to increase speed at its expense, that would also be appreciated. Thank you for your feedback.
import numpy as np
import numba
from numba import njit, prange
@numba.jit(nopython=True, fastmath=True, parallel=True)
def f(A, B):
len_A = A.shape[0]
len_B = B.shape[0]
num_factors = B.shape[1]
C = np.zeros((len_A, len_B))
for i in prange(len_A):
for j in prange(len_B):
for a in prange(num_factors):
A_elem = A[i,a]
B_elem = B[j,a]
AB_elem = (A_elem + B_elem)/2
C[i,j] += A_elem * np.log(A_elem/AB_elem) + \
B_elem * np.log(B_elem/AB_elem) + \
(1-A_elem) * np.log((1-A_elem)/(1-AB_elem)) + \
(1-B_elem) * np.log((1-B_elem)/(1-AB_elem))
C = (np.maximum(C, 0)/2*num_factors)**0.5
return C
#A = np.random.rand(10000, 10000)
#B = np.random.rand(10000, 10000)
#f(A, B)