What is the alternative of numpy.newaxis in tensorflow?

Rahul picture Rahul · Feb 20, 2017 · Viewed 11.6k times · Source

Hi I am new to tensorflow. I want to implement the following python code in tensorflow.

import numpy as np
a = np.array([1,2,3,4,5,6,7,9,0])
print(a) ## [1 2 3 4 5 6 7 9 0]
print(a.shape) ## (9,)
b = a[:, np.newaxis] ### want to write this in tensorflow.
print(b.shape) ## (9,1)

Answer

Divakar picture Divakar · Feb 20, 2017

I think that would be tf.expand_dims -

tf.expand_dims(a, 1) # Or tf.expand_dims(a, -1)

Basically, we list the axis ID where this new axis is to be inserted and the trailing axes/dims are pushed-back.

From the linked docs, here's few examples of expanding dimensions -

# 't' is a tensor of shape [2]
shape(expand_dims(t, 0)) ==> [1, 2]
shape(expand_dims(t, 1)) ==> [2, 1]
shape(expand_dims(t, -1)) ==> [2, 1]

# 't2' is a tensor of shape [2, 3, 5]
shape(expand_dims(t2, 0)) ==> [1, 2, 3, 5]
shape(expand_dims(t2, 2)) ==> [2, 3, 1, 5]
shape(expand_dims(t2, 3)) ==> [2, 3, 5, 1]