CODEWITHSUNDEEP
pythonnumpymultidimensional-arrayarray-broadcasting
H.Rappeport
H.Rappeport

Reputation: 597

Add dimensions to a tensor and duplicate values along the new axes

Suppose I have a 2D ndarray X with 

X.shape == (m, n)

I'd like to add two more dimensions to X, while copying the values along these new axes. i.e. I want 

new_X.shape == (m, n, k, l)

and for all i,j 

new_X[i, j, :, :] = X[i, j]

What's the best way to achieve this?

Upvotes: 1

Views: 46

Answers (1)

Divakar
Divakar

Reputation: 221574

You can simply get a tensor view with np.broadcast_to -

np.broadcast_to(a[...,None,None],a.shape+(k,l)) # a is input array

The benefit would be that there's no extra memory overhead with it and hence virtually free rumtime.

If you need an array output with its own memory space, append with .copy().

Sample run -

In [9]: a =  np.random.rand(2,3)

In [10]: k,l = 4,5

In [11]: np.broadcast_to(a[...,None,None],a.shape+(k,l)).shape
Out[11]: (2, 3, 4, 5)

# Verify memory space sharing
In [12]: np.shares_memory(a,np.broadcast_to(a[...,None,None],a.shape+(k,l)))
Out[12]: True

# Verify virtually free runtime
In [17]: a =  np.random.rand(100,100)
    ...: k,l = 100,100
    ...: %timeit np.broadcast_to(a[...,None,None],a.shape+(k,l))
100000 loops, best of 3: 3.41 µs per loop

Upvotes: 2

Related Questions