Other subsystems

A grab bag of smaller-but-still-meaningful subsystems that don't merit a dedicated page. Each has a one-paragraph orientation and pointers into the code.

torch/lazy/ and torch/_lazy/ — Lazy Tensor Core

A historical experiment that became the foundation of XLA's PyTorch integration: every op queues a lazy IR node; a graph-level optimizer materializes results when forced. Most active usage is via torch_xla (out-of-tree). The in-tree code lives at torch/csrc/lazy/, torch/lazy/, and torch/_lazy/. Codegen for new lazy ops is in torchgen/dest/lazy_ts_lowering.py.

torch/nativert/ — Native runtime

A new C++ runtime for executing AOT-compiled programs without the Python interpreter. Targets the same use cases as TorchScript's static runtime and AOTInductor but with a cleaner foundation. Files at torch/nativert/ (and the corresponding torch/csrc/nativert/).

torch/mtia/ and torch/csrc/mtia/ — Meta MTIA

Hooks for Meta's MTIA accelerator. Mostly device-management Python wrappers; the real kernels are out-of-tree.

torch/xpu/ and torch/csrc/xpu/ — Intel XPU

Hooks for Intel GPUs. Real kernels live in the intel-extension-for-pytorch repo and the in-tree XPU stubs in aten/src/ATen/native/xpu/. C++ runtime helpers in c10/xpu/.

torch/sparse/ and aten/src/ATen/native/sparse/ — Sparse tensors

PyTorch supports COO, CSR, CSC, BSR, BSC sparse layouts. The Python entry point is torch.sparse; ops live in aten/src/ATen/native/sparse/. The dispatch keys Sparse* and SparseCsr* route to these kernels. Active areas: sparse linalg, semi-structured (2:4) sparse, sparse-on-GPU.

torch/package/ — torch.package

Self-contained model packaging: archive a model + its source code dependencies into a single zip that can be loaded later without the original Python environment. Implementation at torch/package/ (Python) and torch/csrc/jit/serialization/ (file format).

torch/utils/_pytree.py — pytrees

JAX-style pytree (nested-Python-container) flattening and mapping. Used pervasively in compile, AOT autograd, distributed, and module APIs. Active contributor: XuehaiPan.

torch/utils/data/ — DataLoader and friends

The dataloader stack: Dataset, DataLoader, Sampler, IterableDataset, multi-worker spawning, pin memory, sharding. The C++ side of multi-worker (queues, watchdog) is in torch/csrc/DataLoader.cpp.

torch/multiprocessing/ — Tensor sharing across processes

PyTorch's torch.multiprocessing shadows stdlib multiprocessing with one twist: tensors sent across processes share storage via cudaIpcMemHandle (CUDA) or shared memory file descriptors (CPU). The shared-memory bookkeeping lives in torch/csrc/CudaIPCTypes.cpp and torch/multiprocessing/reductions.py.

torch/utils/cpp_extension.py — Custom C++ / CUDA extensions

JIT-compile a C++ or CUDA file into a Python extension via torch.utils.cpp_extension.load. Driven by setuptools + Ninja. Active contributors: fmassa, ezyang, malfet.

torch/utils/hipify/ — CUDA → HIP source translation

Source-to-source converter that hipifies the CUDA kernels for ROCm builds. Driven from tools/amd_build/.

torch/utils/checkpoint.py — Activation checkpointing

The classic "save memory by re-running forward during backward" technique, implemented as a torch.autograd.Function plus a saved-tensors-hook variant for selective checkpointing.

torch/distributions/ — Probability distributions

Independent of training: the torch.distributions package implements common probability distributions with sample/rsample/log_prob/entropy and supports reparameterized gradients. Lives entirely in torch/distributions/.

torch/onnx/symbolic_caffe2.py — Caffe2 symbolic export

Vestigial. Most of caffe2/ is in the slow process of being deleted; only the file-format reader/writer (caffe2/serialize/) is still load-bearing for torch.save/torch.jit.save.

torch/cpu/ — CPU runtime helpers

A thin module exposing CPU capability flags (AVX2/AVX512/AMX/NEON detection) and the torch.cpu.amp autocast wrapper. The actual SIMD intrinsics live under aten/src/ATen/cpu/vec/.

torch/backends/ — Backend toggles

torch.backends.{cudnn, cuda, mkl, mkldnn, mps, opt_einsum, openmp, quantized, …} exposes per-backend flags (deterministic algorithms, allow_tf32, fastest convolution algorithm, etc.). Each is a thin wrapper over at::globalContext() getters/setters.

torch/monitor/ — Lightweight monitoring

In-process stats counters and event registry; used by some distributed training stacks for SLO metrics.

tools/ — Build / dev tooling

Beyond the autograd codegen, tools/ has dozens of small utilities: tools/linter/, tools/test/, tools/build_pytorch_libs.py, tools/clang_tidy/, tools/setup_helpers/, tools/amd_build/. Anything that runs during pip install -e . but isn't ATen lives here.

For deeper coverage of any of these, follow the file pointers and read the modules directly. Each is small enough to be self-explanatory once you know it exists.