Open-Source Wikis

/

PyTorch

/

Background

/

Design decisions

pytorch/pytorch

Design decisions

A small set of architectural choices ripple through the entire codebase. Internalizing them makes the rest of PyTorch much easier to read.

"Eager by default"

PyTorch's defining choice is immediate execution of the autograd graph. When you write c = a + b, the addition happens now; no graph is built unless you opt in via torch.compile or torch.export. This is the opposite of TensorFlow 1.x's static-graph default and the foundational reason PyTorch caught on with researchers.

Implications:

  • The dispatcher must be fast — every op call goes through it.
  • Autograd's "tape" is recorded incrementally rather than constructed up front.
  • Performance optimizations (kernel fusion, scheduling) require an additional compilation step rather than being free.

"Tensor is intrusive_ptr<TensorImpl>"

Tensors are reference-counted handles to a TensorImpl. Copy is cheap; ownership is unambiguous; circular references are avoided by weak_intrusive_ptr where needed. This shape:

  • Lets Python and C++ share tensor objects without copying.
  • Lets the dispatcher pass tensors by const Tensor& cheaply.
  • Tightly couples lifetime of the storage and the metadata.

The decision to use intrusive_ptr over std::shared_ptr is for performance — shared_ptr requires a separate control block; intrusive_ptr keeps the refcount inline.

"Storage is separate from Tensor"

A view doesn't allocate. Views and the base share Storage. This makes operations like reshape, slice, transpose, and unsqueeze essentially free — they create a new TensorImpl with different sizes/strides pointing at the same storage. The cost is that some code paths must check is_contiguous() before iterating.

"Dispatcher with key sets"

Every op call goes through c10::Dispatcher::call(...), which picks the highest-priority key in tensor.key_set() | TLS_include_set minus the TLS exclude set. Routing logic lives outside the kernels — the kernels register against keys.

This was a deliberate redesign in 2018-2019. Before it, the kernel-selection logic was a forest of if-else chains. Putting it behind a key-set abstraction made it possible to:

  • Add Autograd, Autocast, Functionalize, FuncTorch, and the Python key as composable "modes" rather than ad-hoc checks.
  • Implement third-party backends (XLA, MTIA, MPS) without touching every op.
  • Express transformations like vmap and grad as kernel registrations.

"Codegen the boilerplate"

Operator declarations live in aten/src/ATen/native/native_functions.yaml; backward formulas live in tools/autograd/derivatives.yaml. Most C++/Python boilerplate (entry-point wrappers, autograd Nodes, Python bindings, type stubs) is generated by torchgen at build time.

This means there is one canonical place to declare an op. The cost is a build dependency on Python code (torchgen) and a layer of indirection that confuses newcomers. The benefit is that adding an op is mostly a YAML change.

"Compile is opt-in, not the default"

torch.compile was added in 2.0 as a function decorator / one-line transformation. It is opt-in: existing code that doesn't call torch.compile keeps running through the eager dispatcher. This makes the compiler stack a strict superset rather than a replacement.

This decision shapes the compiler:

  • Dynamo must handle "real" Python (closures, control flow on tensor values, dynamic types) — it can't require a typed subset.
  • Inductor must produce code that's correct and faster than eager — slower compilation isn't an option.
  • The fallback to eager is a first-class behavior: graph breaks, mode="eager" backend, etc.

"Functions over modules in the compile stack"

Pre-torch.compile, the deployment story was torch.jit.script(MyModule()): a module-level transform. PyTorch 2.0 inverts this — torch.compile(fn) works at the function level. Modules are handled by attaching a compiled forward.

The reason: function-level transforms compose with other function-level transforms (vmap, grad) much more naturally than module-level transforms do.

"Functionalize first"

Both AOTAutograd and torch.export run a functionalization pass that rewrites in-place ops and views into their out-of-place equivalents. After functionalization, the graph is pure and trivially differentiable / serializable / optimizable.

The cost is a more complex tracing pipeline (a FunctionalTensor wrapper, view-replay machinery). The benefit is that downstream passes can be much simpler — they don't have to reason about aliases.

"Composable distributed"

The 2024-era distributed stack is built around two abstractions: DTensor (a sharded Tensor subclass) and DeviceMesh (a group of devices arranged in a multi-dim grid). On top of those:

  • FSDP2 is a per-module wrapper.
  • TP is a per-parameter shard.
  • Pipelining is a per-stage scheduler.

Each is independently usable; they compose because they all operate on DTensor and DeviceMesh. This is the architectural inversion of "DDP / FSDP / TP each had its own wrapper class".

"Subclasses are a public extension point"

torch.Tensor subclasses (with __torch_dispatch__ and friends) are a supported way to define custom tensor semantics — FakeTensor, DTensor, MaskedTensor, NestedTensor, FunctionalTensor, and a long tail of internal subclasses all use this. The Python dispatch key (Python, PythonTLSSnapshot) lets the dispatcher route to user code.

The trade-off: this surface is not ABI-stable, and getting subclass interactions right (with autograd, with views, with the compile stack) is hard.

Built by Factory AutoWiki from public repository content. It is a generated preview for codebase exploration, not source-maintained documentation.

Design decisions – PyTorch wiki | Factory