Custom ops and extensions

What it is

PyTorch supports several ways to extend it with custom kernels:

1. torch.utils.cpp_extension — JIT or setuptools-based C++/CUDA extensions.
2. torch.library — register custom ops at runtime so they participate in the dispatcher, autograd, FakeTensor, and torch.compile.
3. torch.autograd.Function — Python-level custom autograd functions.
4. Out-of-tree backends — register a brand-new backend (e.g., a vendor accelerator) using the PrivateUse1/PrivateUse2/PrivateUse3 dispatch keys.

torch.utils.cpp_extension

The fastest way to ship a single C++/CUDA file. Two flavors:

• JIT load. torch.utils.cpp_extension.load(name, sources=...) compiles on demand at first call.
• Setuptools build. BuildExtension and CppExtension / CUDAExtension integrate with setup.py for prebuilt wheels.

from torch.utils.cpp_extension import load
my_ext = load(name="my_ext", sources=["my_op.cpp", "my_op.cu"], verbose=True)
y = my_ext.my_op(x)

The implementation is in torch/utils/cpp_extension.py. It locates nvcc, picks the right C++ ABI, sets compute capabilities, and forwards to setuptools + Ninja.

torch.library

The user-facing dispatcher API. Lets you:

• Define a brand-new op: torch.library.define("my_lib::my_op", "(Tensor x) -> Tensor").
• Register kernels: @torch.library.impl("my_lib::my_op", "CPU") or "CUDA", "AutogradCPU", etc.
• Register a fake (meta) kernel: @torch.library.register_fake("my_lib::my_op") so tracing works.
• Register an autograd formula: torch.library.register_autograd("my_lib::my_op", ...).

Once registered, the op is indistinguishable from a native ATen op — it shows up in profiling, torch.compile can trace and decompose it, autograd works, and FakeTensor mode can shape-infer it.

import torch

@torch.library.custom_op("my_lib::sin_relu", mutates_args=())
def sin_relu(x: torch.Tensor) -> torch.Tensor:
    return x.sin().relu()

@sin_relu.register_fake
def _(x):
    return torch.empty_like(x)

@sin_relu.register_autograd
def _(ctx, grad):
    (x,) = ctx.saved_tensors
    return grad * x.cos() * (x.sin() > 0).to(grad.dtype)

@sin_relu.register_kernel("Meta")
def _(x):
    return torch.empty_like(x)

The implementation is in torch/library.py and torch/_library/.

torch.autograd.Function

For pure-Python custom backwards. Useful when you don't need a new dispatcher op (e.g., a wrapper around an existing op with a different gradient).

class MyOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.sin()

    @staticmethod
    def backward(ctx, grad_out):
        (x,) = ctx.saved_tensors
        return grad_out * x.cos()

y = MyOp.apply(x)

Pre-PyTorch-2.x this was the main path. It still works, and is required for some advanced patterns (custom backwards over fused ops without a dispatcher op). However, torch.library.custom_op is now preferred because it composes with torch.compile.

Custom backends (PrivateUse1)

For an entire new accelerator there are reserved dispatch keys PrivateUse1, PrivateUse2, PrivateUse3 plus matching AutogradPrivateUse1 etc. Out-of-tree projects (e.g., vendor extensions) typically:

1. Pick PrivateUse1 as their key.
2. Register a name via torch._C._rename_privateuse1_backend("my_accel") so users see device='my_accel'.
3. Register kernels via TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) { ... }.
4. Provide a Python module torch.my_accel mirroring torch.cuda with stream/event/memory APIs.

This is how Intel's intel_extension_for_pytorch (XPU before it became a first-class device) and several internal accelerators integrate.

Where it lives

Path	Contents
torch/utils/cpp_extension.py	load, CppExtension, CUDAExtension
torch/library.py	High-level Python torch.library API
torch/_library/	Internal helpers (custom_op, fake, autograd)
torch/csrc/stable/	The ABI-stable C++ surface for extensions
torch/headeronly/	Header-only types safe for stable-ABI use
torch/header_only_apis.txt	List of header-only APIs
aten/src/ATen/core/Library.h	C++ macro side: TORCH_LIBRARY
torch/autograd/function.py	Function / FunctionCtx

Where to read next

• Systems / Dispatcher — what torch.library plugs into.
• Systems / Autograd — what Function plugs into.
• API / C-API stable — for ABI-stable extensions.