In [31]: x = np.arange(5*12).reshape(5,4,3)
In [32]: x@x.transpose(0,2,1)
Out[32]:
array([[[ 5, 14, 23, 32],
[ 14, 50, 86, 122],
[ 23, 86, 149, 212],
[ 32, 122, 212, 302]],
[[ 509, 626, 743, 860],
[ 626, 770, 914, 1058],
[ 743, 914, 1085, 1256],
[ 860, 1058, 1256, 1454]],
[[ 1877, 2102, 2327, 2552],
[ 2102, 2354, 2606, 2858],
[ 2327, 2606, 2885, 3164],
[ 2552, 2858, 3164, 3470]],
[[ 4109, 4442, 4775, 5108],
[ 4442, 4802, 5162, 5522],
[ 4775, 5162, 5549, 5936],
[ 5108, 5522, 5936, 6350]],
[[ 7205, 7646, 8087, 8528],
[ 7646, 8114, 8582, 9050],
[ 8087, 8582, 9077, 9572],
[ 8528, 9050, 9572, 10094]]])
In [33]: _.shape
Out[33]: (5, 4, 4)
checking
In [34]: x[0,...].dot(x[0,...].T)
Out[34]:
array([[ 5, 14, 23, 32],
[ 14, 50, 86, 122],
[ 23, 86, 149, 212],
[ 32, 122, 212, 302]])
@
or matmul
does a dot product on the last 2 dimensions, while treating the first, as you do, as the 'i' dimension.
x.tanspose(0,2,1)
switches the last 2 dimensions, leaving the first unchanged. For a 2d array x.transpose()
is x.transpose(1,0)
. It is what you expressed as x_i.T
.
As with dot
, the key for a@b
, is that the last dimension of a
works with the 2nd to the last dimension of b
.