I want to calculate softmax/probability using multinomial logit while using longsumexp to avoid overflow. Using numba bring about 2-3x speedup. Can I do better here? Also when I use fastmath=True it does not seem to bring any speedup, so did I do write numba loops in a wrong way?
import numbaimport numpy as npdef get_p_4d(a, lamda): m = a * lamda[:, None][:,None].transpose(0,3,1,2) c = np.max(m, axis=2)[:,None].transpose(0,2,1,3) aa = np.exp(m - c) logsumexp = c + np.log(aa.sum(axis=2)[:,None].transpose(0,2,1,3)) p = np.exp(m - logsumexp) return p@numba.njit()def get_p_4d_nb(a, lamda, num_code, num_draw, num_action): p = np.empty((num_code, num_draw, num_action, 3)) a = a.transpose(0, 1, 3, 2) for i in range(num_code): for j in range(num_draw): this_lamda = lamda[i,j] for k in range(num_action): p[i, j, k, 0] = a[i, j, k, 0] * this_lamda p[i, j, k, 1] = a[i, j, k, 1] * this_lamda p[i, j, k, 2] = a[i, j, k, 2] * this_lamda c = p[i,j,k,0] c = max(c, p[i,j,k,1]) c = max(c, p[i,j,k,2]) logsumexp = np.log( np.exp(p[i, j, k, 0] - c) + np.exp(p[i, j, k, 1] - c) + np.exp(p[i, j, k, 2] - c)) + c p[i, j, k, 0] = np.exp(p[i, j, k, 0] - logsumexp) p[i, j, k, 1] = np.exp(p[i, j, k, 1] - logsumexp) p[i, j, k, 2] = np.exp(p[i, j, k, 2] - logsumexp) return p.transpose(0, 1, 3, 2)a=np.ones((112,1000,3,3))lamda = np.random.uniform(0., 1., size=112*1000).reshape(112,1000)get_p_4d(a, lamda)get_p_4d_nb(a, lamda, 112, 1000, 3)
