Tensor Shape, Rank, dtype and tf.reshape Explained

The Story — Twelve Eggs, Many Cartons

Every tensor has three key facts: its shape (size of each dimension), its rank (how many dimensions), and its dtype (the type of each value). Most TensorFlow errors you will ever see are shape or dtype errors. So this lesson is worth learning well.

Try It — Climb the Rank Ladder

Reading Shape, Rank and Size in Code

import tensorflow as tf

x = tf.zeros([2, 3, 4])

print("shape :", x.shape)          # TensorShape, known in Python
print("ndim  :", x.ndim)           # rank as a Python int
print("rank  :", tf.rank(x))       # rank as a tensor
print("size  :", tf.size(x))       # total number of values
print("dims  :", x.shape[0], x.shape[-1])
print("list  :", x.shape.as_list())

x.shape vs tf.shape(x)

@tf.function(input_signature=[tf.TensorSpec(shape=[None, 3])])
def batch_info(x):
    print("static  :", x.shape)          # printed once at trace time
    tf.print("dynamic :", tf.shape(x))   # printed at run time
    return x

batch_info(tf.zeros([5, 3]))

Axes and Negative Indexes

Axes are numbered from the outside in. For a tensor of shape (2, 3, 4), axis 0 has size 2 and axis 2 has size 4. Negative numbers count from the end. So axis=-1 always means the last axis.

dtype — The Type of Every Value

TensorFlow Does Not Mix dtypes for You

NumPy quietly converts types. TensorFlow does not. It raises an error instead. This protects you from silent bugs.

a = tf.constant([1, 2, 3])          # int32
b = tf.constant([0.5, 0.5, 0.5])    # float32

try:
    a + b
except tf.errors.InvalidArgumentError as e:
    print("Error:", e.message.split(":")[0])

print(tf.cast(a, tf.float32) + b)   # fix: cast first

tf.cast — Change the dtype

prices = tf.constant([2.7, -2.7, 9.99])
print(tf.cast(prices, tf.int32))      # cuts toward zero, no rounding!
print(tf.cast(prices, tf.bool))       # non-zero -> True

pixels = tf.constant([0, 128, 255], dtype=tf.uint8)
print(tf.cast(pixels, tf.float32) / 255.0)   # classic image scaling

Try It — The dtype Caster

tf.reshape — Same Values, New Box

tf.reshape(x, new_shape) keeps all the values in the same order. It only changes how they are grouped. There is one rule: the total number of values must stay the same.

x = tf.range(12)
print(tf.reshape(x, [3, 4]))
print(tf.reshape(x, [2, -1]).shape)       # -1 becomes 6
print(tf.reshape(x, [2, 2, 3]).shape)
print(tf.reshape(x, [-1]).shape)          # flatten to 1-D

try:
    tf.reshape(x, [5, 3])
except tf.errors.InvalidArgumentError as e:
    print("Error:", e.message.split(" [Op")[0])

Row-Major Order

TensorFlow reads values in row-major order. It walks along the last axis first, like reading a book: left to right, then the next line. Reshape pours values into the new shape in that same order.

Try It — The Reshape Animator

Adding and Removing Axes of Size 1

Many layers need an extra axis. For example, a model trained on batches cannot take one single image. You add a batch axis of size 1.

img = tf.zeros([28, 28])                  # one grayscale image

batch = tf.expand_dims(img, axis=0)       # add batch axis at the front
print(batch.shape)

with_channel = img[..., tf.newaxis]       # add channel axis at the end
print(with_channel.shape)

print(tf.squeeze(tf.zeros([1, 28, 28, 1])).shape)   # drop ALL size-1 axes
print(tf.squeeze(tf.zeros([1, 28, 28, 1]), axis=0).shape)

Shapes You Will Meet in Deep Learning

Golden Rules