CODEWITHSUNDEEP
pythonnumpy
Leo
Leo

Reputation: 288

Taking a tensor product in python numpy without performing a sum

I have two tensors, each 2D, with a common long axis (eg 20.000) and diffenrent short axes eg one 9 the other 10). I want to end up with a 9X10X20000 tensor, such that for each location on the long axis, the other two axes are the tensor product. Explicitly, with the "long" axis here 4, I want to do:

A = np.arange(8).reshape(2,4)
B = np.arange(12).reshape(3,4)
C = np.zeros(2,3,4)
for i in range(2):
    for j in range(3):
        for k in range(4):
            C[i,j,k] = A[i,k]*B[j,k]

This code works, but I was wondering: is there a numpy way of doing this, without running for loops?

The context is for training a neural net, with the long axis being the training examples. I get a formula of this form when calculating gradients of the cost function. Cheers, Leo

Upvotes: 2

Views: 1729

Answers (2)

Jaime
Jaime

Reputation: 67457

This particular one you can also get to work with plain old broadcasting:

>>> A[:, np.newaxis, :] * B[np.newaxis, :, :]
array([[[ 0,  1,  4,  9],
        [ 0,  5, 12, 21],
        [ 0,  9, 20, 33]],

       [[ 0,  5, 12, 21],
        [16, 25, 36, 49],
        [32, 45, 60, 77]]])

The above is equivalent to the simpler:

>>> A[:, None] * B
array([[[ 0,  1,  4,  9],
        [ 0,  5, 12, 21],
        [ 0,  9, 20, 33]],

       [[ 0,  5, 12, 21],
        [16, 25, 36, 49],
        [32, 45, 60, 77]]])

Upvotes: 2

DSM
DSM

Reputation: 353419

I tend to use np.einsum for these problems, which makes it easy to specify what should happen in terms of the indices:

>>> A = np.arange(8).reshape(2,4)
>>> B = np.arange(12).reshape(3,4)
>>> np.einsum('ik,jk->ijk', A, B)
array([[[ 0,  1,  4,  9],
        [ 0,  5, 12, 21],
        [ 0,  9, 20, 33]],

       [[ 0,  5, 12, 21],
        [16, 25, 36, 49],
        [32, 45, 60, 77]]])

Upvotes: 3

Related Questions