DOCSEARCHbyAlgolia
Submit Docs

PyTorch Documentation

website·Indexed Jun 18, 2026

https://docs.pytorch.org/docs/stable/

Official documentation for PyTorch, covering core packages like torch.autograd, torch.amp, and torch.cuda. Includes guides on accelerator integration, device management, automatic mixed precision, distributed training with DDP, C++ extensions, and community governance.

Tokens
187.1K
Snippets
284
Records
1.2K
Agent score
50%

What's inside PyTorch

  1. Overview of torch.func (formerly functorch)

    torch.func (previously known as functorch) provides JAX-like composable function transforms for PyTorch. It allows users to apply higher-order functions to numerical functions to compute different quantities.

    Key capabilities include:

    • Auto-differentiation transforms: e.g., grad(f) returns a function that computes the gradient of f.
    • Vectorization/Batching transforms: e.g., vmap(f) returns a function that computes f over batches of inputs.
    • Composition: Transforms can be composed arbitrarily, such as vmap(grad(f)) to compute per-sample gradients, which is otherwise inefficient in standard PyTorch.

    Note: This library is currently in beta. APIs may change, and there may not be full coverage over all PyTorch operations.

  2. Overview of ExponentialFamily base class

    The torch.distributions.exp_family.ExponentialFamily is an abstract base class for probability distributions belonging to an exponential family.

    Distributions inheriting from this class can leverage optimized methods for computing .entropy() and analytic KL divergence using the AD (Automatic Differentiation) framework and Bregman divergences. This is particularly useful for ensuring the correctness of entropy and KL divergence calculations in models using these distribution types.

  3. Overview of PyTorch DTensor (Distributed Tensor)

    PyTorch DTensor (Distributed Tensor) is a torch.Tensor subclass designed for the SPMD (single program, multiple data) programming model. It provides a single-device abstraction for multi-device programming by handling sharded storage, operator computation, and collective communications transparently across a DeviceMesh.

    Key concepts:

    • DeviceMesh: Represents the device topology and communicators of the cluster as an n-dimensional array.
    • Placement: Describes the sharding layout of the logical tensor on the DeviceMesh. Supported types include:
      • Shard: Tensor is sharded on a specific dimension dim across the DeviceMesh dimension.
      • Replicate: Tensor is replicated across the DeviceMesh dimension.
      • Partial: Tensor is pending a reduction operation on the DeviceMesh dimension.

    Note: torch.distributed.tensor is currently in alpha state.

  4. Overview of DeviceMesh abstraction

    DeviceMesh is a high-level abstraction that manages process groups (or NCCL communicators). It simplifies the creation of inter-node and intra-node process groups by handling rank setup automatically based on a provided mesh shape describing the device topology. Use init_device_mesh() to create a new DeviceMesh.

  5. Overview of PyTorch Accelerator Integration Pathway

    Since PyTorch 2.1, a streamlined integration pathway has been established for adding new hardware accelerators to the PyTorch ecosystem. This pathway leverages refinements to the PrivateUse1 Dispatch Key, core subsystem extension mechanisms, and device-agnostic refactoring of modules like torch.accelerator and memory management.

    Following this modern integration path ensures:

    • Speed: Integration can be done independently in downstream codebases without modifying upstream PyTorch code.
    • Future-proofing: New PyTorch features will automatically support the accelerator if this standard path is followed.
    • Autonomy: Vendors can manage their own integration timelines without relying on upstream community review bandwidth.

    The integration surface is categorized into four major axes:

    1. Runtime: Core components like Event, Stream, Memory, Generator, Guard, Hooks, and C++ scaffolding.
    2. Operators: Implementation of forward/backward operators, fallback operators, fallthroughs, and STUBs in C++ and Python.
    3. Python Frontend: Python bindings for modules and device-agnostic APIs.
    4. High-level Modules: Integration with subsystems such as AMP (Automatic Mixed Precision), Compiler, ONNX, and Distributed.

  6. Overview of PyTorch C++ API capabilities

    PyTorch provides a C++ frontend that mirrors much of the Python API functionality. Key capabilities include:

    • Tensor and Autograd: Access to torch::Tensor methods (e.g., add, reshape, clone), a tensor indexing API that behaves like the Python version, and the torch::autograd package for building dynamic neural networks.
    • Model Authoring: Full capability to author and train neural networks purely in C++ using components like torch::nn, torch::nn::functional, and torch::optim.
    • Libtorch: The core library (libtorch) containing these C++ APIs.

  7. Overview of torch.fx components

    torch.fx is a toolkit for transforming nn.Module instances using a three-part pipeline:

    1. Symbolic Tracer: Performs "symbolic execution" by feeding Proxy objects through the code to record operations.
    2. Intermediate Representation (IR): A Graph consisting of a list of Node instances representing inputs (placeholder), operations (get_attr, call_function, call_module, call_method), and return values (output).
    3. Python Code Generation: Converts the Graph into valid Python code. This is encapsulated in a GraphModule, which is a torch.nn.Module containing the Graph and a generated forward method.

  8. Overview of torch.futures.Future

    The torch.futures.Future class encapsulates an asynchronous execution of a callable. It is primarily used by the Distributed RPC Framework to manage asynchronous results. It provides APIs to add callbacks, set results or exceptions, and block until a value is ready.

    Warning: GPU support is currently a beta feature and subject to changes.

  9. Overview of the Elastic Agent in torchelastic

    The Elastic Agent serves as the control plane for torchelastic. It is a process responsible for launching and managing underlying worker processes. Its primary responsibilities include:

    1. Distributed Torch Integration: Providing workers with the necessary information to call torch.distributed.init_process_group() successfully.
    2. Fault Tolerance: Monitoring workers and, upon detecting failures or unhealthiness, tearing down all workers and restarting the group.
    3. Elasticity: Reacting to membership changes by restarting workers with the new membership set.

    Agents can be deployed per node (managing local processes) or can be more advanced, managing workers remotely or operating in a decentralized or coordinated manner.

  10. Initialize a torch.Generator

    The torch.Generator class manages the state of the algorithm that produces pseudo-random numbers. It can be used as a keyword argument in many random sampling functions to ensure reproducibility. You can specify the device (e.g., 'cpu' or 'cuda') upon initialization.
    >>> g_cpu = torch.Generator()
>>> g_cuda = torch.Generator(device='cuda')

  11. Constrain dynamic values as tensor sizes using torch._check

    When exporting a model with torch.export, if a value derived from a tensor (via .item()) is used to define the shape of another tensor, the value might not be known at tracing time. You can use torch._check to provide hints about the value's range. This allows the exporter to trace through shape-dependent operations by treating the value as a symbolic integer with known constraints.

    In the example below, torch._check is used to constrain the integer a between 0 and 5, which is then used as a dimension in torch.zeros.

    # mypy: allow-untyped-defs
import torch

class ConstrainAsSizeExample(torch.nn.Module):
    def forward(self, x):
        a = x.item()
        # Provide hints to the exporter about the range of 'a'
        torch._check(a >= 0)
        torch._check(a <= 5)
        return torch.zeros((a, 5))

example_args = (torch.tensor(4),)
model = ConstrainAsSizeExample()

# The exported program will include range constraints for the symbolic value
exported_program = torch.export.export(model, example_args)

  12. Use torchrun for distributed training

    torchrun (a console script for the torch.distributed.run module) is used to spawn multiple distributed training processes on training nodes. It supports single-node and multi-node training for both CPU and GPU.

    When using GPUs, each process operates on a single GPU from index 0 to nproc_per_node - 1.

    Note on Argument Parsing: Since PyTorch 2.0.0, torchrun passes --local-rank=<rank> to your script. For backward compatibility with older scripts, you should handle both --local-rank and --local_rank in your argparse configuration.

    import argparse
parser = argparse.ArgumentParser()
# Handle both dashed and underscored versions for compatibility
parser.add_argument("--local-rank", "--local_rank", type=int)
args = parser.parse_args()