Let h be a function which applies entrywise (e.g squaring, etc.). Let u and v be a 1d arrays of length k and A be a 2d array of shape n x m. Here l, n, and m can be very large!
Question. In numpy, what is an efficient way to compute the n x m array defined by B := v[0] * h(u[0] * A) + ... + v[k-1] * h(u[k-1] * A) ?
Observations
- Of course, a naive solution is to just loop over the u[j]'s and v[j]'s and accumulate v[j] * h(u[j] * A).
- Another solution which seems horrible (at least memory-wise is): B = np.sum(v[:, None, None] * h(u[:, None, None] * A), axis=0)