CodeVetter · Repo Unpack · Extension points
pytorch/pytorch
PyTorch is built around registries. Operators extend via native_functions.yaml + torchgen (in-tree) or torch.library (Library.define/impl, custom_op, register_fake, register_autograd) and appear lazily under torch.ops.<ns>. Compiler backends plug into torch.compile through the torch_dynamo_backends entry-point group or @register_backend. Distributed backends register via the torch.distributed.backends entry-point group. Device vendors use _register_device_module/Privateuse1. Interop is through __torch_function__ (overrides.py) and __torch_dispatch__ modes (_python_dispatch.py). Serialization accepts custom location tags via register_package, and C++ extensions JIT-build via cpp_extension.
- In-tree op extension: declare `func:` schema + per-backend `dispatch:` entries in native_functions.yaml, implement the kernel in aten/src/ATen/native/*.cpp, and torchgen generates C++/Python bindings — the documented 'modern mechanism for adding operators'.
- torch.library.Library(ns, kind in DEF/IMPL/FRAGMENT, dispatch_key) registers new operator schemas (define()) or overrides kernels for a specific dispatch key (impl()); reserved namespaces are rejected and registration is finalized via weakref cleanup.
- Higher-level Python op APIs exist alongside Library: torch.library.custom_op, register_fake (meta/fake kernels for compile), register_autograd, and impl_abstract — the supported path for making custom ops work under torch.compile.
- torch.ops namespaces are lazily materialized: _OpNamespace.__getattr__ creates OpOverloadPackets on demand, and each packet resolves overloads via torch._C._get_operation_overload — so newly registered ops are reachable as torch.ops.ns.op without code changes.
- Compile backends plug in two ways: the @register_backend decorator populating _BACKENDS/_COMPILER_FNS/_BACKEND_TAGS, or external packages exposing entry points in the 'torch_dynamo_backends' group discoverable by list_backends().
- Distributed backends are entry-point plugins: pyproject registers mpi/gloo/nccl/nccl-legacy/nccl2/nccl-lazy/ucc/xccl/nccl4py factories under group 'torch.distributed.backends', and init_process_group accepts 'third-party plugin' backend strings and device-qualified specs like 'cpu:gloo,cuda:custom_backend'.
- Out-of-tree device runtimes use torch._register_device_module(device_type, module) which installs the module as torch.<device_type> in sys.modules and the torch namespace.
- __torch_function__ overrides are handled in torch/overrides.py (with a non-overridable allowlist and performance warnings about changing this file); __torch_dispatch__ modes live in torch/utils/_python_dispatch.py.
- Serialization is extensible: user extensions register custom location tags plus tagging/deserialization methods via torch.serialization.register_package, and the safe-globals API (add_safe_globals/safe_globals) extends what weights_only load accepts.
- JIT C++/CUDA extension path: torch.utils.cpp_extension.load()/load_inline() write ninja files, compile objects and produce loadable extension modules — the out-of-tree kernel authoring route.
- Inductor's compile_fx accepts a decompositions dict (else select_decomp_table) and config_patches (with per-device '<device>.<key>' routing), making decomposition and config the knob surface for backend behavior.
- AOTInductor C shim is extensible: torchgen's --extend-aoti-c-shim generates c shims for out-of-tree ATen ops from an out-of-tree native_functions.yaml into torch/csrc/inductor/aoti_torch/generated/extend.