TensorFlow
neural networks
graph creation
tensor shape
machine learning

Tracking tensor shape at graph creation time

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

When you build TensorFlow graphs, there are two different shape questions: what TensorFlow knows statically when the graph is created, and what becomes known only when real data flows through the graph. Tracking shape at graph creation time means inspecting or constraining the static shape information before execution starts. That is useful because shape problems caught early are much easier to debug than runtime failures buried deep in a training step.

Static Shape Versus Dynamic Shape

TensorFlow exposes shape information in two major ways.

Static shape is available from the tensor object itself.

python
1import tensorflow as tf
2
3x = tf.keras.Input(shape=(32, 16))
4print(x.shape)

This prints a TensorShape object, which may include known dimensions and unknown ones such as the batch size.

Dynamic shape is computed at runtime:

python
1import tensorflow as tf
2
3x = tf.keras.Input(shape=(32, 16))
4print(tf.shape(x))

tf.shape(x) is a tensor operation, so it belongs to execution time rather than pure graph-construction time.

Inspect Static Shapes While Building Layers

If you want to track shapes during graph creation, inspect tensor.shape or tf.keras.backend.int_shape() while wiring the model.

python
1import tensorflow as tf
2from tensorflow.keras import layers
3
4inputs = tf.keras.Input(shape=(28, 28, 1))
5print("input:", inputs.shape)
6
7x = layers.Conv2D(16, 3, padding="same")(inputs)
8print("after conv:", x.shape)
9
10x = layers.MaxPooling2D()(x)
11print("after pool:", x.shape)
12
13x = layers.Flatten()(x)
14print("after flatten:", x.shape)

This kind of tracing is simple and very effective when you are debugging shape mismatches in a new architecture.

Enforce Expected Shapes Early

Inspection is useful, but constraints are even better. If you know what shape a tensor should have, tell TensorFlow explicitly.

python
1import tensorflow as tf
2
3@tf.function
4def use_tensor(x):
5    x = tf.ensure_shape(x, [None, 32])
6    return x * 2

tf.ensure_shape tells TensorFlow what shape is expected and raises an error if the actual shape is incompatible.

You can also tighten shapes with set_shape() when building pipelines:

python
1import tensorflow as tf
2
3x = tf.keras.Input(shape=(None, 16))
4print(x.shape)
5
6x.set_shape([None, 10, 16])
7print(x.shape)

Use this only when you genuinely know the shape constraint is valid.

TensorSpec Helps With Function Signatures

When you decorate a function with @tf.function, TensorSpec can document the expected rank and dimensions at trace time.

python
1import tensorflow as tf
2
3@tf.function(input_signature=[tf.TensorSpec(shape=[None, 20], dtype=tf.float32)])
4def project(x):
5    return tf.matmul(x, tf.ones((20, 4)))

This does two things:

  • it documents the expected input shape clearly
  • it reduces accidental retracing caused by inconsistent input structure

That makes graph creation more predictable.

Why Shape Tracking Matters In Practice

Many TensorFlow shape errors come from a mismatch that could have been seen earlier:

  • flattening the wrong rank
  • concatenating tensors with incompatible last dimensions
  • feeding a time-distributed layer data with missing sequence axes
  • assuming a fixed sequence length when the model was built for variable length

By checking static shapes as you build the graph, you move the error closer to its cause.

Common Pitfalls

The most common mistake is treating tf.shape(tensor) as if it were a compile-time answer. It is a runtime tensor, not a pure static shape description.

Another issue is relying on partially known shapes without noticing that some dimensions are still None. That can make a model look more constrained than it really is.

It is also easy to overuse set_shape() and force an invalid assumption into the graph. If the actual data does not satisfy the constraint, later failures can become harder to reason about.

Finally, remember that Keras model summaries and layer output shapes are often the easiest first debugging tool. Do not skip them in favor of more complicated inspection logic.

Summary

  • Graph-creation-time shape tracking is about static shape information, not runtime tensor values.
  • Use tensor.shape and int_shape() to inspect shapes while building the graph.
  • Use tf.ensure_shape or TensorSpec when you want TensorFlow to enforce expectations early.
  • Distinguish clearly between static shapes and tf.shape(...) runtime values.
  • Catching shape issues during graph construction makes TensorFlow models much easier to debug.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.