Dot product between rows of two matrices with same batch dimension

user_4561314 0
Viewed 180

Suppose I have two 2-dimensional tensors with the same batch dimension (that is, same number of rows).

import numpy as np

# Both t1 and t2 have shape (2, 3)

t1 = np.array([[1.0, 2.0, 3.0], [10.0, 11.0, 12.0]])
# array([[ 1.,  2.,  3.],
#        [10., 11., 12.]])

t2 = np.array([[21.0, 22.0, 23.0], [30.0, 31.0, 32.0]])
# array([[21., 22., 23.],
#        [30., 31., 32.]])

How do I compute the dot product between the rows of the two tensors? Specifically, I want to end up with the desired (2, 1) tensor result:

desired_result = np.array([np.dot(t1[0], t2[0]), np.dot(t1[1], t2[1])])
# array([ 134., 1025.])

I already tried:

np.dot(t1, t2)
# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
#   File "<__array_function__ internals>", line 5, in dot
# ValueError: shapes (2,3) and (2,3) not aligned: 3 (dim 1) != 2 (dim 0)

np.tensordot(t1, t2)
# array(1159.)

Thanks for any help.

numpy python
3 Answers

Since both t1 and t2 are 2d arrays, numpy.dot does matrix multiplication. For your case, you can just multiply the two arrays element wise and then sum rows:

(t1 * t2).sum(axis=1)
# [ 134. 1025.]
user_49834500

Using einsum:

np.einsum("ij,ij->i", t1, t2)
user_149232270

np.matmul/@ does batched dot, but the batch dimension is first of 3 (or more). So we have to expand the arrays:

In [36]: t1[:,None,:]@t2[:,:,None]
Out[36]: 
array([[[ 134.]],

       [[1025.]]])
In [37]: _.shape
Out[37]: (2, 1, 1)

and then squeeze out the extra dimension(s)

matmul is fast, but for this task the other answers might be clearer.

user_9019250
Related