mlsanityData to 7 Oct 2026

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

arXiv 2112.08429 ↗First posted 15 Dec 2021cs.LGcited by 38

Not on the map. Why?

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.

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

Builds on

23 of 50 references.

References · 23

  1. 2017Attention Is All You Needcited by 59,932
  2. 2015Deep Residual Learning for Image Recognitioncited by 51,809
  3. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
  4. 2014Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translationcited by 4,016
  5. 2017Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inferencecited by 1,548
  6. 2017In-Datacenter Performance Analysis of a Tensor Processing Unitcited by 610
  7. 2018Quantizing deep convolutional networks for efficient inference: A whitepapercited by 507
  8. 2019Deep Learning Recommendation Model for Personalization and Recommendation Systemscited by 327
  9. 2014Caffe: Convolutional Architecture for Fast Feature Embeddingcited by 287
  10. 2015MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systemscited by 195
  11. 2019Learning to Paint With Model-based Deep Reinforcement Learningcited by 70
  12. 2016Theano: A Python framework for fast computation of mathematical expressionscited by 64
11 more
  1. 2020MLIR: A Compiler Infrastructure for the End of Moore's Lawcited by 64
  2. 2018Don't Unroll Adjoint: Differentiating SSA-Form Programscited by 44
  3. 2021FBGEMM: Enabling High-Performance Low-Precision Deep Learning Inferencecited by 40
  4. 2019TensorFlow Eager: A Multi-Stage, Python-Embedded DSL for Machine Learningcited by 22
  5. 2019Chainer: A Deep Learning Framework for Accelerating the Research Cyclecited by 22
  6. 2017DyNet: The Dynamic Neural Network Toolkitcited by 21
  7. 2021LazyTensor: combining eager execution with domain-specific compilerscited by 11
  8. 2017Training Deep AutoEncoders for Collaborative Filteringcited by 10
  9. 2018AutoGraph: Imperative-style Coding with Graph-based Performancecited by 9
  10. 2021Swift for TensorFlow: A portable, flexible platform for deep learningcited by 6
  11. 2018Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUscited by 3

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2022: 1; 2023: 11; 2024: 5; 2025: 6; 2026: 15.

38 papers.

Sort:

By area:

  1. 2024Wukong: Towards a Scaling Law for Large-Scale Recommendationcited by 72
  2. 2025PyG 2.0: Scalable Learning on Real World Graphscited by 49
  3. 2023Oobleck: Resilient Distributed Training of Large Models Using Pipeline Templatescited by 40
  4. 2024Allo: A Programming Model for Composable Accelerator Designcited by 16
  5. 2023Relax: Composable Abstractions for End-to-End Dynamic Machine Learningcited by 12
  6. 2022Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Modelscited by 11
  7. 2023Slapo: A Schedule Language for Progressive Optimization of Large Deep Learning Model Trainingcited by 11
  8. 2023Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Modelscited by 5
  9. 2024Integrated Hardware Architecture and Device Placement Searchcited by 4
  10. 2025Quasar: A Programming Language Specialized for LLM Code Actionscited by 4
  11. 2023LinguaLinked: A Distributed Large Language Model Inference System for Mobile Devicescited by 3
  12. 2023Efficient ConvBN Blocks for Transfer Learning and Beyondcited by 2
26 more
  1. 2023A Dataflow Compiler for Efficient LLM Inference using Custom Microscaling Formatscited by 2
  2. 2024Enhanced Sparsification via Stimulative Trainingcited by 2
  3. 2026ExecuTorch -- A Unified PyTorch Solution to Run AI Models On-Devicecited by 2
  4. 2026KernelSight-LM: A Kernel-Level LLM Inference Simulatorcited by 2
  5. 2026LEMUR 2: Unlocking Neural Network Diversity for AIcited by 2
  6. 2026GraphIR: Architecture-Level Search States for LLM-Guided Neural Architecture Evolutioncited by 2
  7. 2025Collapsing Taylor Mode Automatic Differentiationcited by 1
  8. 2025GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2cited by 1
  9. 2025Scalable Synthesis of distributed LLM workloads through Symbolic Tensor Graphscited by 1
  10. 2026Self-Distillation Policy Optimization via Visual Feedback: Bridging Code and Visual Artifactscited by 1
  11. 2023CMLCompiler: A Unified Compiler for Classical Machine Learningcited by 0
  12. 2023MGit: A Model Versioning and Management Systemcited by 0
  13. 2023DiviML: A Module-based Heuristic for Mapping Neural Networks onto Heterogeneous Platformscited by 0
  14. 2023A Survey on Design Methodologies for Accelerating Deep Learning on Heterogeneous Architecturescited by 0
  15. 2024Kernel Approximation using Analog In-Memory Computingcited by 0
  16. 2025Flashlight: PyTorch Compiler Extensions to Accelerate Attention Variantscited by 0
  17. 2026What Artificial Intelligence can do for High-Performance Computing systems?cited by 0
  18. 2026Accelerating Vision Transformers on Brain Processing Unitcited by 0
  19. 2026NEST: Network- and Memory-Aware Device Placement For Distributed Deep Learningcited by 0
  20. 2026Forge-UGC: FX optimization and register-graph engine for universal graph compilercited by 0
  21. 2026AgentCompile: An LLM-Guided Compiler for Direct CUDA Inferencecited by 0
  22. 2026XPR: An Extensible Cross-Platform Point-Based Differentiable Renderercited by 0
  23. 2026Efficient EEG Seizure Detection Using INT8 Quantization, Channel Pruning, and Spiking Neural Networkscited by 0
  24. 2026Memory-Efficient Activation Checkpointing with Sliding Window and Hirschberg's Algorithm for 0/1 Knapsack Solving in PyTorchcited by 0
  25. 2026CIERA: Cross-Iteration Exponent Reuse for Lossless Allgather in Sharded MoE Trainingcited by 0
  26. 2026Trident: Unifying Guarded Dispatch and Host Execution for PyTorch Triton Workloadscited by 0