pytorch/pytorch
github.com/pytorch/pytorch ↗ · 3995 files · 36.9 MB · scan capped
-------------------------------------------------------------------------------- PyTorch is a Python package that provides two high-level features: - Tensor computation (like NumPy) with strong GPU acceleration - Deep neural networks built on a tape-based auto
Overview
PyTorch is a hybrid Python/C++ deep-learning framework shipping a GPU-accelerated tensor library, a tape-based autograd engine, and the torch.compile JIT stack (Dynamo bytecode tracing -> AOT Autograd -> Inductor/Triton codegen). The Python package `torch` is a thin but enormous facade (torch/__init__.py is ~3600 lines) over the compiled `torch._C` extension, which loads libtorch via ctypes and dispatches operators through a C++ dispatcher whose schemas are declared in aten/src/ATen/native/native_functions.yaml and turned into bindings by the torchgen code generator. Since this commit the build has migrated to scikit-build-core (PEP 517): pyproject.toml is the real manifest, CMakeLists.txt orchestrates the native build, and setup.py is a deprecated forwarding shim scheduled for removal in PyTorch 2.18.
System map
PyTorch is a layered system: the c10 core library and ATen tensor library (C++), the autograd engine in torch/csrc/autograd, and the Python frontend package `torch` bound through torch._C (torch_python). Build is PEP 517 via scikit-build-core: pyproject.toml declares the backend, CMakeLists.txt is the native build entry point, and cmake/Codegen.cmake runs torchgen over native_functions.yaml to emit ATen headers/sources plus autograd and Python binding sources. Operator dispatch crosses the Python/C++ boundary through OpOverloadPacket/OpOverload in torch/_ops.py into the C++ dispatcher. CI is a large GitHub Actions matrix of reusable _binary-build/_binary-test workflows gated by a job-filter workflow; test selection uses target determination. Packaging installs CMake artifacts under the `torch` wheel directory with an extensive exclusion list to keep the wheel lean.
- The build backend is scikit-build-core (build-backend = "scikit_build_core.build"); PEP 517 builds run CMake, and requires-python is >=3.10.
- setup.py is a deprecated shim: `install`/`develop` forward to `pip install --no-build-isolation`, all other commands fail, and the file is scheduled for removal in PyTorch 2.18.
- CMakeLists.txt is the native build entry point driven by scikit-build-core; its header comment documents the ordered stages: EnvVarForwarding, PreBuildSteps (git submodules + pinned NCCL fetch), Dependencies.cmake, codegen via cmake/Codegen.cmake (torchgen over native_functions.yaml) and tools/setup_helpers/generate_code.py, then add_subdirectory of torch/headeronly, c10 and caffe2.
- `import torch` loads native deps before Python bindings: _load_global_deps() ctypes.CDLLs libtorch_global_deps.so with RTLD_GLOBAL, inspects /proc/self/maps for libcudart, then _preload_cuda_deps preloads cublasLt/cublas/cudnn/nvrtc/cudart/cupti/nccl etc. from site-packages before `from torch._C import *`.
- On Windows, _load_dll_libraries() adds DLL directories, preloads vcruntime/msvcp DLLs, and LoadLibraryExW-loads every *.dll under torch/lib before importing torch._C.
- torchgen/gen.py main() is the codegen CLI invoked at build time: it parses aten/src/ATen/native/native_functions.yaml and native/tags.yaml and emits ATen headers/sources, dispatcher registrations, and the AOTInductor C shim, with flags like --rocm/--mps/--xpu/--mtia and --per-operator-headers.
- The wheel ships packages torch, torchgen and functorch, with CMake artifacts installed into install-dir "torch"; a long exclude list strips static archives (*.a), C++ sources, third-party cmake configs, docs and dev files from the wheel.
- Version and runtime dependencies are dynamic metadata: tools/metadata version provider resolves PYTORCH_BUILD_VERSION/PKG-INFO/version.txt, and the dependencies provider injects CUDA/ROCm/XPU pip deps from PYTORCH_EXTRA_INSTALL_REQUIREMENTS.
- Console scripts are torchrun (-> torch.distributed.run:main) and torchfrtrace (-> flight recorder trace); the torch.distributed.backends entry-point group registers mpi/gloo/nccl/ucc/xccl/nccl4py backend factories.
- The Python op surface is assembled at import: _C._initExtension(torch_shm_manager path) runs first, then every callable in _C._VariableFunctions is copied into torch's globals (appending to __all__), followed by submodule imports (cpu, cuda, xpu, mps, mtia, distributed, nn, optim, fx-related) and _C._init_names.
- Native dependencies are vendored as git submodules under third_party/ (pybind11, googletest, protobuf, gloo, NNPACK, pthreadpool, benchmark, and more).
- CI is GitHub Actions: trunk.yml runs on pushes to main/release/*/landchecks plus a nightly cron, fronted by job-filter.yml and target_determination.yml (with llm_td_retrieval) that gate which reusable _binary-build/_binary-test workflows execute; pull.yml does the same for PRs.
Feature catalog
User-facing surface is huge: eager tensor ops bound from _VariableFunctions, the nn.Module/optim training stack, tape-based autograd, torch.compile (Dynamo+Inductor), torch.export and legacy TorchScript for graph capture, torch.distributed with pluggable collectives backends plus the torchrun elastic launcher, checkpointing via torch.save/load, data loading and multiprocessing helpers, quantization/sparse/fft/linalg/profiler subpackages, and C++ extension authoring via cpp_extension. New features typically surface either as torch.* functions (codegen-bound aten ops) or as subpackage imports in torch/__init__.py.
- Eager tensor ops are exposed by copying every member of _C._VariableFunctions into the torch namespace at import (PRIVATE_OPS like unique_dim are skipped; __all__ is extended for non-underscore names).
- torch.nn.Module is the model base class providing parameters(), state_dict()/load_state_dict(), and the __call__ -> forward dispatch used by all neural network code.
- torch.autograd.backward() and torch.autograd.grad() drive the tape-based engine: backward validates inputs, handles __torch_function__ overrides, builds grad tensors via _make_grads, then calls _engine_run_backward which invokes Variable._execution_engine.run_backward in C++.
- torch.compile(model, fullgraph=, backend=, mode=, options=, dynamic=, recompile_limit=, isolate_recompiles=, dynamic_shapes=) is the JIT entry point: it wraps the backend string (default inductor) in _TorchCompileInductorWrapper/_TorchCompileWrapper and returns torch._dynamo.optimize(...)(model).
- compile modes are "default", "reduce-overhead" (CUDA graphs), "max-autotune" and "max-autotune-no-cudagraphs"; options keys include epilogue_fusion, max_autotune, fallback_random, shape_padding, triton.cudagraphs, trace.enabled and guard_filter_fn.
- torch.export.export()/save()/load() provide AoT graph capture producing ExportedPrograms; torch.compile is explicitly a no-op inside an export region.
- TorchScript lives on via torch.jit (torch/jit/_script.py) and mobile .ptl artifacts; SECURITY.md treats them as executable code requiring trust.
- torch.distributed.init_process_group supports gloo/nccl/mpi/ucc/xccl and third-party backends, auto-selecting nccl for cuda, gloo for cpu, xccl for xpu when backend is omitted; collectives run over a Store established by rendezvous or init_method URL.
- torch.save writes a zip container (records: data.pkl, .format_version, .storage_alignment, byteorder, data/<key> storages); torch.load reads it with weights_only=True by default since 2.6, supporting map_location, mmap and a safe-globals allowlist.
- Accelerator and utility subpackages imported at the end of __init__: cpu, cuda, mps, mtia, xpu, amp, autocast, distributions, fft, linalg, nested, sparse, special, profiler, hub, futures, multiprocessing, optim, nn, ao (quantization), quasirandom, masked.
- Higher-order and functional transforms (vmap, cond, sym_* operators) are part of the public API listed in torch.__all__.
- Custom C++/CUDA extensions are built JIT via torch.utils.cpp_extension.load()/load_inline() which generate ninja build files and compile into loadable modules.
Data flow
The core data path is: Python call -> torch.* binding (codegen'd _VariableFunctions) or torch.ops namespace -> C++ dispatcher keyed by DispatchKey -> backend kernel (CPU/CUDA/etc.) -> TensorImpl/Storage in c10/ATen, with autograd Nodes attached during forward when grad is required. Compile mode inserts a bytecode interception layer: Dynamo's PEP 523 hook traces Python frames into FX graphs, AOT Autograd splits forward/backward, and Inductor lowers to Triton/C++ kernels. Persistence flows through a custom zip format where storages are extracted via pickle persistent_id and written as raw byte records. Distributed state flows through Store-based rendezvous into backend ProcessGroups.
- torch.ops.<ns>.<op> lazily resolves through _OpNamespace.__getattr__ -> OpOverloadPacket.__getattr__ -> torch._C._get_operation_overload + _get_schema -> cached OpOverload; calling it invokes self._op (the C++ unboxed dispatcher entry).
- OpOverload keeps a per-DispatchKey _dispatch_cache and a py_kernels table; decompose() prefers a Python CompositeImplicitAutograd kernel and otherwise redispatches into C++ via _op_dk.
- native_functions.yaml is the schema source of truth: each entry is `func: name(ArgType...) -> Return` plus a `dispatch:` block mapping backend keys (CPU, CUDA, ...) to kernel function names; torchgen parse_native_yaml turns these into NativeFunctions + backend indices that drive all generated bindings.
- Autograd data flow: forward ops record grad_fn Nodes; backward() -> _engine_run_backward stashes the contextvars context in TLS then calls Variable._execution_engine.run_backward; C++ Engine::execute validates outputs, creates a GraphTask and schedules NodeTasks on per-device worker threads.
- torch.save path: save() -> _save() runs a PyTorchPickler whose persistent_id intercepts storages, dedupes them by _cdata into serialized_storages, enforces dtype-consistency per data_ptr, writes data.pkl + .format_version + .storage_alignment + byteorder records, then streams each storage (moved to CPU, optionally via pinned-memory D2H) as data/<key>.
- torch.load path: stricter _is_zipfile check -> _open_zipfile_reader (torch._C.PyTorchFileReader) -> _load() reads byteorder/.format_version/.storage_alignment, computes restore_location from map_location, and load_tensor() builds UntypedStorages (meta storages when a fake mode is active, recording _checkpoint_offset for layout-preserving reload).
- weights_only=True is the default when pickle_module is unset; _weights_only_unpickler restricts pickle GLOBALs to an allowlist extensible via add_safe_globals()/safe_globals(), and get_unsafe_globals_in_checkpoint scans data.pkl for non-allowlisted globals.
- torch.compile data flow: compile() -> torch._dynamo.optimize -> _optimize -> OptimizeContext wrapping the fn; the C++ PEP 523 eval-frame hook intercepts frames, symbolic_convert builds an FX graph in OutputGraph with guards, then the backend (compile_fx) is invoked and generated bytecode calls the compiled artifact with resume functions after graph breaks.
- Inductor's compile_fx orchestrates end-to-end compilation: selects the decomposition table, wakes the AsyncCompile worker pool, wraps inner_compile into fw_compiler/bw_compiler callables, and calls aot_autograd which partitions the graph and invokes those compilers before fx_codegen_and_compile emits runnable code.
- init_process_group resolves the default backend (Backend(...) validation, device-based default), performs rendezvous through init_method/env:// to get a Store, then _new_process_group_helper constructs the ProcessGroup; all collective payloads traverse the backend with no auth or encryption.
- Compile caches are per code object, not per frame: multiple compiled entries accumulate per code object until torch._dynamo.config.recompile_limit (default 8), after which execution falls back to eager; isolate_recompiles gives a compile call its own entry set.
- Distributed checkpoint (torch.distributed.checkpoint) deliberately bypasses the weights_only protections: checkpoints are assumed to come from the user's own trusted store, so loading may execute arbitrary code.
Behavior traces
Traced flows: (1) import torch walks native-library preloading before any Python API exists; (2) torch.compile goes Python wrapper -> dynamo OptimizeContext -> C++ frame hook -> FX graph -> inductor -> aot_autograd -> Triton codegen -> guard-cached bytecode; (3) torch.save/load round-trips through a pickle-metadata + raw-storage zip format with the restricted unpickler on the read side; (4) backward() descends from torch.autograd into the C++ Engine GraphTask scheduler; (5) init_process_group does store rendezvous then backend construction; (6) the test runner shards test files and invokes pytest per module.
- import torch order: _rocm_init.initialize() if present -> (Windows) _load_dll_libraries / (POSIX) _load_global_deps ctypes.CDLL(libtorch_global_deps, RTLD_GLOBAL) + optional _preload_cuda_deps -> `from torch._C import *` -> _C._initExtension(torch_shm_manager path) -> copy _C._VariableFunctions into globals -> import subpackages (cpu/cuda/xpu/mps/mtia before multiprocessing) -> _C._init_names(storage classes) -> attach docstrings -> bind torch.ops/torch.classes namespaces.
- torch.compile call sequence: compile() logs API usage, rejects Python >=3.16 / free-threaded <3.13.3, resolves backend via get_default_backend, wraps inductor in _TorchCompileInductorWrapper, then torch._dynamo.optimize -> _optimize -> check_if_dynamo_supported -> (TORCHDYNAMO_DISABLE/justknobs check) -> optimize_assert if nopython else _optimize_catch_errors(convert_frame.convert_frame(backend,...)) -> OptimizeContext; per-call the C++ eval-frame hook runs guards then either executes cached compiled code or triggers compilation.
- Inside a dynamo compile: InstructionTranslator in symbolic_convert.py executes bytecode into VariableTrackers; OutputGraph owns the FX graph/guards; compile_subgraph calls the backend; PyCodegen emits replacement bytecode; resume_execution.py builds continuation functions for graph breaks; CheckFunctionManager builds a C++ GuardManager tree for fast recheck.
- torch.save sequence: save() -> _open_zipfile_writer -> _save() -> PyTorchPickler(persistent_id) dumps obj to data.pkl -> write .format_version/.storage_alignment/byteorder records -> for each serialized storage, move to CPU (pinned-memory D2H if configured) -> zip_file.write_record('data/<key>', storage, num_bytes).
- torch.load sequence: torch._C._log_api_usage_once('torch.load') -> default weights_only=True -> _is_zipfile -> _open_zipfile_reader -> _load(): read byteorder/format_version/alignment -> Unpickler (weights-only restricted unpickler or full pickle module) resolves 'storage' persistent ids via load_tensor -> restore_location applies map_location -> storage rebuilt from data/<key> record (or lazily via mmap).
- loss.backward() sequence: torch.autograd.backward -> functorch guard -> grad_variables deprecation shim -> __torch_function__ override check -> _tensor_or_tensors_to_tuple + _make_grads -> retain_graph defaults to create_graph -> _engine_run_backward (TLS context stash, optional logging hooks) -> Variable._execution_engine.run_backward -> C++ Engine::execute -> validate_outputs -> GraphTask scheduled on worker threads -> .grad accumulated on leaves (accumulate_grad=True).
- init_process_group sequence: reject double init (GroupMember.WORLD check) -> set_pytorch_distributed_envs_from_justknobs -> Backend(backend) validation/defaulting -> store resolution via explicit store or rendezvous(init_method, default env://) -> _new_process_group_helper builds the backend ProcessGroup (eager ncclCommInit* when device_id is passed for NCCL).
- Test execution: python test/run_test.py -> main() -> parse_args + TESTS_TO_INCLUDE env merge -> get_selected_tests (TEST groups like INDUCTOR_TESTS/DISTRIBUTED_TESTS/DYNAMO_CORE_TESTS) -> target-determination amend_tests -> do_sharding -> run_tests splits serial vs parallel files -> run_test_module invokes pytest per file with -m marker expressions; each test file uses run_tests() from common_utils which can spawn per-test-case subprocesses with a 15-min timeout and one retry.
- Codegen flow at build: cmake/Codegen.cmake runs `python -m torchgen.gen` -> main() parses native_functions.yaml + tags.yaml from aten/src/ATen -> get_custom_build_selector -> emits ATen headers/sources, dispatcher registrations, and the AOTInductor C shim into build/aten/src/ATen and torch/csrc/inductor/aoti_torch/generated.
Testing signals
Tests live in test/ (~411 files) and run under pytest driven by the custom harness test/run_test.py, with unittest-style TestCase classes from torch.testing._internal.common_utils. Device-generic coverage comes from instantiate_device_type_tests, and snapshot assertions from expecttest (assertExpectedInline). CI runs pull.yml on PRs and trunk.yml on main, with a job-filter and target-determination pipeline that shards tests and can run only ~25% via TD. Env vars (PYTORCH_TEST_WITH_*) gate ASAN/ROCM/slow/dynamo/inductor configurations; disabled and slow test lists are consumed from JSON files. C++ tests are built as gtest binaries and discovered by run_test.py.
- pytest.ini configures the suite: testpaths=test, --tb=native, --capture=sys, no warnings plugin, plain asserts, xfail_strict=True, and custom markers serial/multigpu/periodic (sharding via a custom pytest_shard_custom.py, not pytest-shard).
- test/run_test.py partitions TESTS into groups (INDUCTOR_TESTS, DISTRIBUTED_TESTS, DYNAMO_CORE_TESTS, CPYTHON_TESTS, ONNX_TESTS, QUANTIZATION_TESTS, FUNCTORCH_TESTS, CPP_TESTS) and supports TESTS_TO_INCLUDE env, --include/--exclude, sharding (do_sharding) using per-file/class historical times, and TD mode that runs ~25% of tests.
- The shared base is TestCase(expecttest.TestCase) in common_utils (line ~3720) with assertExpectedInline for golden-string snapshot tests; run_tests() lints that every test class subclasses TestCase, supports --discover, --save-xml, subprocess-per-test-case mode, slow/disabled test JSON import, and 15-minute timeouts with one retry.
- Test gating is env-var driven: PYTORCH_TEST_WITH_ASAN/TSAN/UBSAN/ROCM/MTIA/PERIODIC/SLOW/DYNAMO/INDUCTOR/AOT_EAGER, PYTORCH_TEST_SKIP_FAST, PYTORCH_TESTING_DEVICE_EXCEPT_FOR, PYTORCH_TEST_CUDA_MEM_LEAK_CHECK etc. select configurations at module import.
- AGENTS.md mandates the test idiom: subclass torch.testing._internal.common_utils.TestCase, call run_tests() in __main__, use assertEqual for tensor equality, @parametrize for input matrices, and instantiate_device_type_tests for device-generic numeric tests.
- Dynamo tests live in test/dynamo/ using torch._dynamo.test_case.TestCase (auto-resets dynamo and patches strict config); debugging workflows center on TORCH_LOGS=guards,graph_breaks,+dynamo and CompileCounter/EagerAndRecordGraphs helpers.
- CI orchestration: pull.yml (PRs) and trunk.yml (main/release pushes + cron) both call job-filter.yml then target_determination.yml (fed by llm_td_retrieval.yml) before the reusable _binary-build-* and _binary-test-* per-platform workflows.
- Result plumbing: run_test.py uploads artifacts mid-run in CI (parse_xml_and_upload_json + zip_and_upload_artifacts), writes into test/test-reports/, and maps negative exit codes to signal names for crash reporting.
- C++ tests are gtest binaries discovered via _is_cpp_test; googletest is a third_party submodule, and conftest.py/pytest_shard_custom.py are copied into CPP_TESTS_DIR so cpp tests share the pytest sharding path.
- Lint/type gates are configured in-tree: .lintrunner.toml (lintrunner -a before commits), ruff/isort/usort config in pyproject.toml (line-length 88), mypy.ini + mypy-strict.ini and pyrefly.toml.
- Disabled/flaky tests are data-driven: PYTORCH_TEST_RERUN_DISABLED_TESTS re-runs entries from DISABLED_TESTS_FILE, and CI downloads test-case configs and slow-test classifications (SLOW_TESTS_FILE) before running.
Risk map
The dominant risks are deserialization and the trust boundary of model artifacts: torch.load is a pickle surface (mitigated by the now-default weights_only mode), TorchScript/mobile artifacts are executable code, and distributed primitives have zero authentication or encryption by design. The project's own SECURITY.md explicitly scopes out crashes/OOB from invalid inputs and cache poisoning, meaning large parts of the C++ kernel surface are intentionally unhardened. Internally, torch/__init__.py is a fragile monolith with import-order constraints and environment mutation (PATH, dlopen flags), serialization depends on vendored miniz internals, and the sprawling CI matrix plus content-hashed Docker images create operational blast radius.
- torch.load with weights_only=False executes arbitrary pickle — documented as 'Never load data from an untrusted source'; weights_only=True became the default in 2.6 and uses the restricted _weights_only_unpickler with a safe-globals allowlist.
- torch.distributed primitives (c10d, RPC, TCPStore) include no authorization and send unencrypted messages: 'anybody with access to the network can execute arbitrary code with the privileges of the user running PyTorch.'
- Distributed checkpoints intentionally bypass weights_only protections — loading a checkpoint 'may execute arbitrary code'; the trust model assumes your own store produced the artifact.
- TorchScript and mobile .ptl models are treated as untrusted executable code; even inspection tools like torch.utils.model_dump can execute embedded code.
- PyTorch's stated policy is to not validate all inputs: wrong shapes/dtypes/devices can cause crashes or out-of-bounds memory access and are classified as bugs, not vulnerabilities — a deliberate, large unhardened C++ surface.
- Local caches (compiled kernels, Triton cache) are trusted by design; cache poisoning is explicitly out of scope, so any code path writing/reading these caches assumes a non-hostile filesystem.
- CI supply-chain mitigations are documented: sccache caching is susceptible to corruption (artifacts must not run near secrets), non-ephemeral runners risk reverse shells (hence no repo secrets), and release builds require ephemeral runners + cold-start builds + deployment environments.
- torch/__init__.py is a 3588-line module with hard ordering constraints (device modules must precede multiprocessing; quantized/ao imports must be last), global mutation (sys.setdlopenflags, os.environ['PATH'] on Windows, os.add_dll_directory), and a sorted-__all__ runtime assertion — edits here have global blast radius.
- _load's fast path hardcodes miniz internals — data descriptor sizes 24/16 bytes, 0xFFFFFFFF boundary, header-offset arithmetic that 'must be kept in sync with that of miniz' — coupling the Python loader to vendored zip-writer behavior.
- Windows loading mutates process-global state: prepends dll_paths to os.environ['PATH'] and LoadLibraryW's every *.dll in torch/lib — classic DLL search-order/preload surface.
- Generated-file traps: .pyi stubs are generated from .pyi.in templates (edit the template, not the output), and torchgen emits the ATen/op registration surface at build time — editing generated files or native_functions.yaml without understanding codegen silently changes public API.
- The .ci/docker directory is content-hashed: any file change — including the README — forces a full Docker image rebuild for CI; AGENTS.md warns against touching it (e.g., during upstream Ubuntu outages).
Extension points
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.
Contributor handoff
AGENTS.md is mandatory reading and unusually specific: build ONLY via `pip install -e . -v --no-build-isolation`, test via run_tests/TestCase + instantiate_device_type_tests, lint via spin lint/fixlint plus `lintrunner -a` before commits, and never `noqa: S101` an assert (rewrite as raise AssertionError). Edit .pyi.in templates not generated .pyi; never touch .ci/docker (content-hash rebuilds); respect the ghstack workflow (detached HEAD is intentional, never push, preserve trailers). Style: minimal comments, ASCII-only in new comments, explicit state, match existing patterns. Use agent_space/ for scratch. AI_POLICY.md requires marking all AI-generated content in GitHub interactions.
- Build rule: 'All build (both codegen, C++ and python) is done via pip install -e . -v --no-build-isolation. You should NEVER run any other command to build PyTorch.'
- Test rule: tests subclass torch.testing._internal.common_utils TestCase and call run_tests(); assertEqual for tensor equality, @parametrize for input sets, instantiate_device_type_tests for device-generic numerics tests.
- Lint rule: only spin-provided commands (spin lint / spin fixlint); run `lintrunner -a` before committing and fix what it reports.
- S101 rule: never silence 'use of assert' with noqa — rewrite `assert cond, msg` as `if not cond: raise AssertionError(msg)` because python -O strips asserts; do not invert float comparisons (NaN).
- Generated-file rule: many .pyi files are generated from .pyi.in templates — always edit the .pyi.in, never the generated .pyi.
- .ci/docker/ is content-hashed for image rebuilds — any change (even README) triggers a full Docker rebuild; don't touch it unless intending to rebuild, and avoid it entirely when Docker builds are broken upstream.
- Git/ghstack rules: detached HEAD means ghstack flow (commit directly, don't branch); never amend or push ghstack commits unless asked; preserve ghstack-source-id and Pull-Request trailers; submit with ghstack (--no-stack for single commits); never push to gh/USER/N branches.
- Commit message rules: include a Test Plan section with literal fenced commands, explain root cause for bug fixes, disclose AI authorship, and don't bullet-list individual changes.
- Use agent_space/ (git-ignored, repo root) for scratch scripts and throwaway experiments; never commit from it. Check for a .venv in repo root/parent before assuming missing tools.
- Dynamo debugging kit: TORCH_LOGS=guards,graph_breaks,graph_code,+dynamo,bytecode; TORCH_TRACE=dir + tlparse for production artifacts via trace_structured; torch._dynamo.config.patch for config toggles; bytecode_debugger for generated-code segfaults.
- Style rules: minimal self-documenting code, no trivial one-use helpers, explicit class state (no dynamic setattr/getattr), match existing patterns, ASCII-only in new comments, and prefer keeping long lines readable via locals instead of ruff-driven splits (noqa: B950 on the closing quote for golden strings).
- Build configuration is env-var driven (MAX_JOBS, DEBUG, REL_WITH_DEB_INFO, USE_*/BUILD_*); cmake/EnvVarForwarding.cmake is the single source of truth for env vars — add and document new knobs there.
Languages
| Language | Files | Share of source |
|---|---|---|
| Python | 4942 | 49% |
| C | 2534 | 25% |
| C++ | 2333 | 23% |
| Markdown | 512 | docs/config |
| YAML | 218 | docs/config |
| Shell | 177 | 2% |
| JSON | 30 | docs/config |
| Java | 21 | 0% |
Start reading here
Probable source entrypoint: aten/ — inferred from the directory tree and manifest, not the upstream docs.
Source roots
aten/ (412 files) c10/ (412 files) torch/ (410 files)
Repository structure
| Top-level path | Files |
|---|---|
aten/ | 412 |
c10/ | 412 |
benchmarks/ | 412 |
test/ | 411 |
torch/ | 410 |
tools/ | 403 |
docs/ | 347 |
.github/ | 311 |
Docs found: AGENTS.mdCLAUDE.mdCONTRIBUTING.mdREADME.mddocs/README.mddocs/cpp/README.mddocs/cpp/source/api/aten/accessors.mddocs/cpp/source/api/aten/creation.mddocs/cpp/source/api/aten/index.mddocs/cpp/source/api/aten/indexing.md
Package
pyproject.toml — torch, 4 declared dependencies.
optreeopt-einsumpyyaml#
CI & config entrypoints
README.md— Repository readmeAGENTS.md— Agent instructionsCLAUDE.md— Claude instructions.github/workflows/_binary-build-flash-attention-wheel-linux.yml— GitHub Actions workflow.github/workflows/_binary-build-flash-attention-wheel-windows.yml— GitHub Actions workflow.github/workflows/_binary-build-linux.yml— GitHub Actions workflow.github/workflows/_binary-test-linux.yml— GitHub Actions workflow.github/workflows/_binary-test-macos.yml— GitHub Actions workflow.github/workflows/_binary-test-rocm-linux.yml— GitHub Actions workflow.github/workflows/_binary-test-xpu-linux.yml— GitHub Actions workflow
Structural health signals
Average score 5.6/10 across 96 analyzed source files; 96 hotspots surfaced.
scripts/pr_review/test_workflow_contract.py— 1096 lines, score 4.7 (Large file, Deep nesting)benchmarks/dynamo/runner.py— 1379 lines, score 4.7 (Large file, Deep nesting)tools/nightly.py— 1216 lines, score 4.7 (Large file, Deep nesting)
These are deterministic review leads — size, churn, and structural signals — not proof of a bug in this project.
Recent history
5711f1a4586c2026-09-22 — [distributed] Add DeviceMesh.abort() API (#189655)a402343453f82026-09-22 — [dynamo] Forward `name` parameter to `_TorchCompileWrapper` for custom backends (#196229)6a895d1264e42026-09-22 — Support cross stream sync de-allocation safety for accelerators (#197187)b8834a2d75b22026-09-21 — Dynamo 3.15 support (#178393)a08518fdf4e02026-09-21 — Add docker image for python 3.15 (#179094)