# Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python

arXiv: https://arxiv.org/abs/2112.08429 · HTML: https://mlsanity.com/paper/2112.08429/

## Key facts

- Cited by 40 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 50 references that are in the map.
- First posted on arXiv 15 Dec 2021.

## Builds on (most cited first, 23 of 23)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://mlsanity.com/paper/1712.05877/) (2017)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Quantizing deep convolutional networks for efficient inference: A whitepaper](https://mlsanity.com/paper/1806.08342/) (2018)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [Caffe: Convolutional Architecture for Fast Feature Embedding](https://mlsanity.com/paper/1408.5093/) (2014)
- [MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems](https://mlsanity.com/paper/1512.01274/) (2015)
- [Learning to Paint With Model-based Deep Reinforcement Learning](https://mlsanity.com/paper/1903.04411/) (2019)
- [MLIR: A Compiler Infrastructure for the End of Moore's Law](https://mlsanity.com/paper/2002.11054/) (2020)
- [Theano: A Python framework for fast computation of mathematical expressions](https://mlsanity.com/paper/1605.02688/) (2016)
- [Don't Unroll Adjoint: Differentiating SSA-Form Programs](https://mlsanity.com/paper/1810.07951/) (2018)
- [FBGEMM: Enabling High-Performance Low-Precision Deep Learning Inference](https://mlsanity.com/paper/2101.05615/) (2021)
- [TensorFlow Eager: A Multi-Stage, Python-Embedded DSL for Machine Learning](https://mlsanity.com/paper/1903.01855/) (2019)
- [Chainer: A Deep Learning Framework for Accelerating the Research Cycle](https://mlsanity.com/paper/1908.00213/) (2019)
- [DyNet: The Dynamic Neural Network Toolkit](https://mlsanity.com/paper/1701.03980/) (2017)
- [LazyTensor: combining eager execution with domain-specific compilers](https://mlsanity.com/paper/2102.13267/) (2021)
- [Training Deep AutoEncoders for Collaborative Filtering](https://mlsanity.com/paper/1708.01715/) (2017)
- [AutoGraph: Imperative-style Coding with Graph-based Performance](https://mlsanity.com/paper/1810.08061/) (2018)
- [Swift for TensorFlow: A portable, flexible platform for deep learning](https://mlsanity.com/paper/2102.13243/) (2021)
- [Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs](https://mlsanity.com/paper/1810.09868/) (2018)

## Built on by (most cited first, 30 of 40)

- [DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence](https://mlsanity.com/paper/2606.19348/) (2026)
- [Wukong: Towards a Scaling Law for Large-Scale Recommendation](https://mlsanity.com/paper/2403.02545/) (2024)
- [PyG 2.0: Scalable Learning on Real World Graphs](https://mlsanity.com/paper/2507.16991/) (2025)
- [Oobleck: Resilient Distributed Training of Large Models Using Pipeline Templates](https://mlsanity.com/paper/2309.08125/) (2023)
- [Allo: A Programming Model for Composable Accelerator Design](https://mlsanity.com/paper/2404.04815/) (2024)
- [Relax: Composable Abstractions for End-to-End Dynamic Machine Learning](https://mlsanity.com/paper/2311.02103/) (2023)
- [Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Models](https://mlsanity.com/paper/2206.04959/) (2022)
- [Slapo: A Schedule Language for Progressive Optimization of Large Deep Learning Model Training](https://mlsanity.com/paper/2302.08005/) (2023)
- [Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models](https://mlsanity.com/paper/2302.02599/) (2023)
- [ExecuTorch -- A Unified PyTorch Solution to Run AI Models On-Device](https://mlsanity.com/paper/2605.08195/) (2026)
- [Integrated Hardware Architecture and Device Placement Search](https://mlsanity.com/paper/2407.13143/) (2024)
- [Quasar: A Programming Language Specialized for LLM Code Actions](https://mlsanity.com/paper/2506.12202/) (2025)
- [LinguaLinked: A Distributed Large Language Model Inference System for Mobile Devices](https://mlsanity.com/paper/2312.00388/) (2023)
- [Efficient ConvBN Blocks for Transfer Learning and Beyond](https://mlsanity.com/paper/2305.11624/) (2023)
- [A Dataflow Compiler for Efficient LLM Inference using Custom Microscaling Formats](https://mlsanity.com/paper/2307.15517/) (2023)
- [Enhanced Sparsification via Stimulative Training](https://mlsanity.com/paper/2403.06417/) (2024)
- [MusaCoder: Native GPU Kernel Generation with Full-Stack Training on Moore Threads GPU](https://mlsanity.com/paper/2606.04847/) (2026)
- [KernelSight-LM: A Kernel-Level LLM Inference Simulator](https://mlsanity.com/paper/2606.28565/) (2026)
- [LEMUR 2: Unlocking Neural Network Diversity for AI](https://mlsanity.com/paper/2607.06839/) (2026)
- [GraphIR: Architecture-Level Search States for LLM-Guided Neural Architecture Evolution](https://mlsanity.com/paper/2608.01633/) (2026)
- [Collapsing Taylor Mode Automatic Differentiation](https://mlsanity.com/paper/2505.13644/) (2025)
- [GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2](https://mlsanity.com/paper/2509.16248/) (2025)
- [Scalable Synthesis of distributed LLM workloads through Symbolic Tensor Graphs](https://mlsanity.com/paper/2511.10480/) (2025)
- [Self-Distillation Policy Optimization via Visual Feedback: Bridging Code and Visual Artifacts](https://mlsanity.com/paper/2606.10334/) (2026)
- [CMLCompiler: A Unified Compiler for Classical Machine Learning](https://mlsanity.com/paper/2301.13441/) (2023)
- [MGit: A Model Versioning and Management System](https://mlsanity.com/paper/2307.07507/) (2023)
- [DiviML: A Module-based Heuristic for Mapping Neural Networks onto Heterogeneous Platforms](https://mlsanity.com/paper/2308.00127/) (2023)
- [A Survey on Design Methodologies for Accelerating Deep Learning on Heterogeneous Architectures](https://mlsanity.com/paper/2311.17815/) (2023)
- [Kernel Approximation using Analog In-Memory Computing](https://mlsanity.com/paper/2411.03375/) (2024)
- [Flashlight: PyTorch Compiler Extensions to Accelerate Attention Variants](https://mlsanity.com/paper/2511.02043/) (2025)

## Abstract (arXiv)

Modern deep learning frameworks provide imperative, eager execution programming interfaces embedded in Python to provide a productive development experience. However, deep learning practitioners sometimes need to capture and transform program structure for performance optimization, visualization, analysis, and hardware integration. We study the different designs for program capture and transformation used in deep learning. By designing for typical deep learning use cases rather than long tail ones, it is possible to create a simpler framework for program capture and transformation. We apply this principle in torch.fx, a program capture and transformation library for PyTorch written entirely in Python and optimized for high developer productivity by ML practitioners. We present case studies showing how torch.fx enables workflows previously inaccessible in the PyTorch ecosystem.