tf.newaxis operation in TensorFlow

Question

x_train = x_train[..., tf.newaxis].astype("float32")

x_test = x_test[..., tf.newaxis].astype("float32")

Can someone please explain how tf.newaxis works ?

I found a brief mention in the documentation

https://www.tensorflow.org/api_docs/python/tf/strided_slice

but I could not properly understand.

Answer 1

Check this example:

a = tf.constant([100])
print(a.shape) ## (1)
expanded_1 = tf.expand_dims(a,axis=1)
print(expanded_1.shape) ## (1,1)
expanded_2 = a[:, tf.newaxis]
print(expanded_2.shape) ## (1,1)

It is similar to expand_dims() which adds a new axis.

If you want to add a new axis at the beginning of the tensor, use

expanded_2 = a[tf.newaxis, :]

otherwise (at the end)

expanded_2 = a[:,tf.newaxis]

Answer 2

You can also add dimensions to your tensor whilst keeping the same information present using tf.newaxis.

# Create a rank 2 tensor (2 dimensions)
rank_2_tensor = tf.constant([[10, 7],
                             [3, 4]])

print("dimension: ", rank_2_tensor.ndim)
print("shape    : ", rank_2_tensor.shape)

output:

dimension: 2
shape: TensorShape([2, 2])

# Add an extra dimension (to the end)
rank_3_tensor = rank_2_tensor[..., tf.newaxis] 
# in Python "..." means "all dimensions prior to"

print("dimension: ", rank_3_tensor .ndim)
print("shape    : ", rank_3_tensor .shape)

output:

dimension: 3
shape: TensorShape([2, 2, 1])

You can achieve the same using tf.expand_dims().

rank_new_3_tensor = tf.expand_dims(rank_2_tensor, axis=-1) # "-1" means last axis
print("dimension: ", rank_new_3_tensor .ndim)
print("shape    : ", rank_new_3_tensor .shape)

output:

dimension: 3
shape: TensorShape([2, 2, 1])