Tensor subclasses

Purpose

PyTorch supports user-defined tensor subclasses — Python classes that inherit from torch.Tensor and intercept __torch_function__ and/or __torch_dispatch__ to customize behaviour. The same machinery is used internally by some of the most important features:

• FakeTensor — shape/dtype-only tensors used by Dynamo and AOT Autograd.
• FunctionalTensor — Python wrapper for the C++ functionalize key.
• NestedTensor — ragged tensors (jagged batches of variable-length sequences).
• MaskedTensor — value + boolean mask pairs.
• TwoTensor / CommTensor / debug subclasses — used by tests and tracing infrastructure.

This page is the tour of the subclass machinery and the most-used in-tree subclasses.

Directory layout

Path	Contents
torch/_subclasses/	The core in-tree subclasses
torch/_subclasses/fake_tensor.py	FakeTensor
torch/_subclasses/fake_utils.py	FakeTensor helpers
torch/_subclasses/functional_tensor.py	FunctionalTensor
torch/_subclasses/meta_utils.py	Meta-tensor utilities used by FakeTensor
torch/nested/	NestedTensor
torch/masked/	MaskedTensor
torch/utils/_python_dispatch.py	The TorchDispatchMode API
torch/overrides.py	The __torch_function__ machinery
aten/src/ATen/FunctionalTensorWrapper.h/.cpp	C++ functionalize key (separate from FunctionalTensor)
aten/src/ATen/NestedTensorImpl.h/.cpp	C++ nested tensor impl

Key abstractions

Type	File	Purpose
__torch_function__	torch/overrides.py	Python hook: intercept all torch.* and method calls
__torch_dispatch__	torch/utils/_python_dispatch.py	Python hook: intercept dispatcher calls
TorchDispatchMode	torch/utils/_python_dispatch.py	Push a mode onto the dispatch stack
FakeTensor / FakeTensorMode	torch/_subclasses/fake_tensor.py	Shape-only tensors
FunctionalTensor / FunctionalTensorMode	torch/_subclasses/functional_tensor.py	Python wrapper for functionalize key
NestedTensor	torch/nested/_internal/nested_tensor.py	Ragged batches
MaskedTensor	torch/masked/maskedtensor/core.py	Value + mask

How it works

There are two interception points for tensor subclasses:

__torch_function__

A class method that intercepts every torch.* function and Tensor.method call before dispatch. Used for "outer" customization — e.g., a tensor that auto-converts dtype, a logging wrapper, a CUDA event tracker. The implementation in torch/overrides.py walks the args, finds the highest-priority subclass with a __torch_function__, and calls it.

__torch_dispatch__

A class method that intercepts every dispatcher call after autograd. This is the "inner" hook — it sees the same ATen ops the dispatcher sees. Almost all serious tensor subclasses (FakeTensor, FunctionalTensor, NestedTensor) use it. The implementation pushes a Python dispatch key onto the keyset; the dispatcher routes to a Python-side fallback that calls __torch_dispatch__.

TorchDispatchMode lets you push a global interceptor that catches every op call inside a with block, regardless of input subclass — used by make_fx, FakeTensorMode, and FunctionalTensorMode.

graph TD
    Call[Python: x.add(y)]
    Call --> TF{has __torch_function__?}
    TF -- yes --> TFCall[call __torch_function__]
    TF -- no --> Disp[Dispatcher]
    Disp --> AG[Autograd kernels]
    AG --> Mode{TorchDispatchMode<br/>active?}
    Mode -- yes --> Md[mode.__torch_dispatch__]
    Mode -- no --> Sub{has __torch_dispatch__?}
    Sub -- yes --> SC[subclass.__torch_dispatch__]
    Sub -- no --> BE[Backend kernel]

FakeTensor

A FakeTensor carries shape, dtype, device, stride but no real storage. FakeTensorMode is a TorchDispatchMode that intercepts every op call and consults the corresponding meta kernel (registered in torch/_meta_registrations.py and via aten Meta: dispatch entries) to compute output shape/dtype without doing real compute.

This is foundational for torch.compile: every graph trace runs under FakeTensorMode so it can move through Python without touching the GPU.

FunctionalTensor

FunctionalTensorMode pushes the Functionalize C++ key for tracing-friendly evaluation: in-place ops become out-of-place, view-then-mutate becomes a clone, and aliasing is normalized. AOT autograd uses this on top of FakeTensorMode to produce pure functional FX graphs.

NestedTensor

A NestedTensor (torch/nested/) is a sequence of variable-length tensors stored as a single contiguous buffer plus a [batch + 1] offsets vector. torch/csrc/nested/ and aten/src/ATen/native/nested/ implement the kernels. Used heavily for transformer-attention masks and variable-length sequences.

The "jagged" layout (torch.jagged) is the modern incarnation, integrated with torch.compile and DTensor.

MaskedTensor

A pair (values, mask) that propagates mask through ops. The implementation lives at torch/masked/maskedtensor/; experimental but useful for "ignore these entries" semantics.

Integration points

• Dynamo / AOT Autograd / Inductor all assume __torch_dispatch__ works for FakeTensor.
• Functionalize key kernels in C++ (FunctionalTensorWrapper.cpp) and the Python FunctionalTensorMode interoperate.
• Quantization PT2E uses subclasses to mark ops for quantization annotation.
• FlexAttention uses subclasses to capture user "score modifier" functions.

Entry points for modification

• New tensor subclass → subclass torch.Tensor and implement __torch_function__ and/or __torch_dispatch__. For dispatch-mode behaviour use TorchDispatchMode.
• New meta kernel → register via torch.library.impl(..., "Meta") or in torch/_meta_registrations.py.
• For docs and patterns, see torch/utils/_python_dispatch.py (well-commented) and the FakeTensor implementation.

Key source files

File	Purpose
torch/utils/_python_dispatch.py	TorchDispatchMode
torch/overrides.py	__torch_function__ plumbing
torch/_subclasses/fake_tensor.py	FakeTensor
torch/_subclasses/functional_tensor.py	FunctionalTensor
torch/_meta_registrations.py	Meta kernels
torch/nested/_internal/nested_tensor.py	NestedTensor
aten/src/ATen/FunctionalTensorWrapper.cpp	C++ functionalize
aten/src/ATen/NestedTensorImpl.cpp	C++ nested tensor