How to vectorize custom algorithms in numpy or pytorch?
Suppose I have two matrices:

A: size k x m

B: size m x n


Using a custom operation, my output will be k x n.

This custom operation is not a dot product between the rows of A and columns of B. Suppose this custom operation is defined as:

For the Ith row of A and Jth column of B, the i,j element of the output is:

sum( (a[i] + b[j]) ^20 ), i loop over I, j loops over J


The only way I can see to implement this is to expand this equation, calculate each term, them sum them.

Is there a way in numpy or pytorch to do this without expanding the equation?

Apart from the method @hpaulj outlines in the comments, you can also use the fact that what you are calculating is essentially a pair-wise Minkowski distance:

import numpy as np
from scipy.spatial.distance import cdist

k,m,n = 10,20,30
A = np.random.random((k,m))
B = np.random.random((m,n))

method1 = ((A[...,None]+B)**20).sum(axis=1)
method2 = cdist(A,-B.T,'m',p=20)**20

np.allclose(method1,method2)
# True

