CODEWITHSUNDEEP
pythonarraysnumpy
a_student
a_student

Reputation: 203

How can I apply np.tensordot 'along an axis', e.g. one axis that indicates samples?

Consider two multidimensional arrays, where the first axis in each indicates for example different samples, such that both arrays have the same dimension on this axis, and we want to apply np.tensordot along this dimension between those arrays.

That is, we would like to achieve the following:

import numpy as np
a = np.random.normal(size=(10,2,4,5))
b = np.random.normal(size=(10,2))
out = np.zeros((10,4,5))

for i in range(10):
    out[i] = np.tensordot(b[i],a[i],axes=1)

But using numpy operations/indexing.

Is there a smart way to do this (also with different dimensionalities, these were just examples)

Upvotes: 1

Views: 26

Answers (1)

yann ziselman
yann ziselman

Reputation: 2002

I think the einsum function can help you with that

import numpy as np
a = np.random.normal(size=(10,2,4,5))
b = np.random.normal(size=(10,2))
out = np.zeros((10,4,5))

for i in range(10):
    out[i] = np.tensordot(b[i],a[i],axes=1)

out1 = np.einsum('ijkl,ij->ikl', a, b)

(out1 == out).all()

output:

True

Upvotes: 1

Related Questions