Control Flow Operators in PyTorch

Representing control flows in machine learning (ML) compilers and intermediate representations (IRs) has been a long-standing problem, with many ML compilers opting to avoid supporting them altogether. A common practice is to trace the model and get a straight-line computational graph, for example, by specializing if predicates or by unrolling loops. Although this strategy may work well in some cases, it can result in performance issues and long compilation times. Furthermore, it is problematic in cases where the control flow of the program depends on data or a dynamic shape.In this paper, we present the PyTorch control flow operator library, which addresses the challenge through the introduction of five control flow operators to PyTorch. We will explain their usage, present real world use cases and demonstrate their benefits such as the ability to capture dynamic control flows in the PyTorch 2 IR with cond and reduce compilation time and peak memory usage when rewriting loops with map, scan, associativescan and whileloop.