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 len_B = B.shape num_factors = B.shape 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)