CODEWITHSUNDEEP
pythonnumpy
Maybe
Maybe

Reputation: 2279

The simplest way to expand a tensor so that it has the same dimension as another tensor

Let flag be a one dimensional vector. Let o1 and o2 be some arrays of more than one dimension; maybe 2 or 4 or something else.

Is there a simpler way to achieve the following thing in numpy?

flag = np.random.randint(0, 2, 10)
# we take o1 and o2 as a four dimensional tensor for an example
o1 = np.ones((10, 2, 4, 4))
o2 = np.zeros((10, 2, 4, 4))
while len(flag.shape) < len(o1.shape):
    flag = np.expand_dims(flag, -1)
o = np.where(flag, o1, o2)
print(o)

Upvotes: 0

Views: 363

Answers (1)

Kevin
Kevin

Reputation: 3348

I think what you have looks fine, you could remove the while loop and add the broadcasted dimensions analytically. This can either be done as my comment or you could use np.reshape, which is slightly more readable.

flag = np.random.randint(0, 2, 10)
o1 = np.ones((10, 2, 4, 4))
o2 = np.zeros((10, 2, 4, 4))
diff = o1.ndim - flag.ndim
flag = flag.reshape(-1, *(1,)*diff)
o = np.where(flag, o1, o2)

You calculate the difference in dimensions between flag and o1. Then add that many dimensions at the end of flag as empty dimensions. A tuple multiplied with a scalar with repeat it diff times and it is then being unpacked as an argument to np.reshape.

To address my comment, the slice notation : is for selecting multiple indices within a range. If you leave it empty it will select all indices and this is equivalent to slice(None, None, None) or slice(None). slice works much like range in regards to its parameter and i am basically doing the same thing as explained above.

Upvotes: 1

Related Questions