Google I/O returns May 18-20! Reserve space and build your schedule Register now

tf.expand_dims

TensorFlow 1 version View source on GitHub

Returns a tensor with an additional dimension inserted at index axis.

Given a tensor input, this operation inserts a dimension of size 1 at the dimension index axis of input's shape. The dimension index axis starts at zero; if you specify a negative number for axis it is counted backward from the end.

This operation is useful if you want to add a batch dimension to a single element. For example, if you have a single image of shape [height, width, channels], you can make it a batch of one image with expand_dims(image, 0), which will make the shape [1, height, width, channels].

Examples:

t = [[1, 2, 3],[4, 5, 6]] # shape [2, 3]
tf.expand_dims(t, 0)
<tf.Tensor: shape=(1, 2, 3), dtype=int32, numpy=
array([[[1, 2, 3],
        [4, 5, 6]]], dtype=int32)>
tf.expand_dims(t, 1)
<tf.Tensor: shape=(2, 1, 3), dtype=int32, numpy=
array([[[1, 2, 3]],
       [[4, 5, 6]]], dtype=int32)>
tf.expand_dims(t, 2)
<tf.Tensor: shape=(2, 3, 1), dtype=int32, numpy=
array([[[1],
        [2],
        [3]],
       [[4],
        [5],
        [6]]], dtype=int32)>
tf.expand_dims(t, -1) # Last dimension index. In this case, same as 2.
<tf.Tensor: shape=(2, 3, 1), dtype=int32, numpy=
array([[[1],
        [2],
        [3]],
       [[4],
        [5],
        [6]]], dtype=int32)>

This operation is related to:

  • tf.squeeze, which removes dimensions of size 1.
  • tf.reshape, which provides more flexible reshaping capability

input A Tensor.
axis Integer specifying the dimension index at which to expand the shape of input. Given an input of D dimensions, axis must be in range [-(D+1), D] (inclusive).
name Optional string. The name of the output Tensor.

A tensor with the same data as input, with an additional dimension inserted at the index specified by axis.

ValueError If axis is not specified.
InvalidArgumentError If axis is out of range [-(D+1), D].