I often find myself to broadcast 1d arrays into a specific dimension dim_a given a total number of dimension dim_total. What I mean is the following:

import numpy as np

a = np.arange(10)

dim_a = 2
dim_total = 4
shape = tuple([-1 if idx == dim_a else 1 for idx in range(dim_total)]) 
print(a.reshape(shape))


axis = list(range(dim_total))
del axis[dim_a]
print(np.expand_dims(a, axis=axis))

Both work as expected, however the question is whether there is an even shorter way to achieve this for a single array?

I often find myself to broadcast 1d arrays into a specific dimension dim_a given a total number of dimension dim_total. What I mean is the following:

import numpy as np

a = np.arange(10)

dim_a = 2
dim_total = 4
shape = tuple([-1 if idx == dim_a else 1 for idx in range(dim_total)]) 
print(a.reshape(shape))


axis = list(range(dim_total))
del axis[dim_a]
print(np.expand_dims(a, axis=axis))

Both work as expected, however the question is whether there is an even shorter way to achieve this for a single array?

• python
• arrays
• numpy
• array-broadcasting

• Hide the code in a function. Then you won't care about the 'length'. Look at the expand_dims source code to see how it creates the reshape parameter. – hpaulj Commented Mar 14 at 21:48
• Yes, make sense. I was already close to open a feature request to Numpy, but wanted to see if there are obvious solutions I have missed. The one proposed is already really nice! – Axel Donath Commented Mar 14 at 22:12
• It's easy to reshape, but a bit more work to use your particular mix of inputs. – hpaulj Commented Mar 15 at 0:55
• So you want to expand a to a (1,1,10,1) shape. For many broadcasting operations it would be sufficient to expand it to (10,1), with the leading dimensions added automatically as needed. – hpaulj Commented Mar 15 at 19:50
• Fair point. I think I have a slight personal preference for keeping the (empty) leading dimensions as well and sticking with a single coordinate axes convention for all arrays involved. But this is probably a matter of taste. I also often use keepdims=True and only on the final result squeeze the empty axes. – Axel Donath Commented Mar 15 at 20:42

2 Answers 2

Shorter way to get shape:

shape, shape[dim_a] = [1] * dim_total, -1

Attempt This Online!

Though I'd prefer it in two lines:

shape = [1] * dim_total
shape[dim_a] = -1

I just investigated some more myself and found this solution making use of the walrus operator:

import numpy as np

a = np.arange(10)

dim_a, dim_total = 2, 4

(shape := [1] * dim_total)[dim_a] = -1
np.reshape(a, shape)

I like that it's very compact, but the := is still not very commonly used.