How PyTorch Handles Dynamic Tensor Shapes

AI models must handle environments where input sizes keep varying, like prompt length, image batches, or video clips. This kind of dynamic tensor shape is a major obstacle to performance optimization, but PyTorch 2's torch.compile handles it practically.

PyTorch's default execution mode, eager mode, performs operations one at a time. On GPUs or accelerators like Furiosa's RNGD, the round-trip overhead between Python and hardware grows large, so PyTorch executes most operations asynchronously to hide that cost.

But hiding overhead is different from eliminating it. So PyTorch 2 tries to transform the program into a more compiler-friendly dataflow graph, optimizing and executing it in larger chunks. To do this, it reads Python bytecode and traces the model by following the functions it calls.

In this process, fake tensors play a key role. By executing code with tensors that hold no actual data, it records operations and tensor shapes, building an FX graph without using memory. As in the example, x * 2 + 3 gets transformed into an FX graph with a concrete shape like f32[8], which then becomes the input for further optimization or kernel generation.

The problem is that the first compilation is usually done assuming a static shape. When a different shape comes in afterward, PyTorch needs to verify that the same code can run safely, so it recompiles while marking that tensor as dynamic shape. The remaining tensors can still stay static, making only the necessary parts flexible.

At this point, PyTorch tracks shape constraints symbolically. For example, addition requires the shapes of two tensors to be equal or broadcastable, and matmul requires the contracting dimensions to match. Based on these conditions, it creates guards, reducing the logical expressions with SymPy and even generating C++ check functions.

In the example, the guard creates a condition requiring the input length to be at least 2 or greater, and in the matmul example, it restricts the allowed range along with requiring the two dimensions to match. In other words, compiled code doesn't blindly accept all shapes—it's only reused within a range that has actually been proven safe.

Dynamic shape compilation is still largely a research topic. In practice, the following strategies matter.

• padding to fit specific bucket sizes
• keeping a vectorized fast path alongside a slow path to handle leftovers
• prioritizing optimization for frequently occurring shapes, like profile-guided optimization
• running rare shapes in eager mode to save on compilation cost

Ultimately, the key point is that PyTorch isn't trying to eliminate shape changes themselves, but rather manages that variability through constraints, recompilation, and guards to secure both performance and flexibility together. Dynamic shapes are tricky, but with proper compiler support, they can be handled practically even in large-scale deployment environments.