Numpy Matrix multiplication Maximum in Linear time

Question

I have two big arrays A is n x k and B is k x m.

I would like to compute np.argmax(A @ B, axis=1), th