PyTorch

https://docs.pytorch.org/docs/2.12/
Admin
PyTorch is an optimized tensor library for deep learning using GPUs and CPUs, providing machine...
Tokens:2,040,698
Snippets:17,459
Trust Score:10
Update:2 months ago
Show doc for...
Tokens:
Raw
### Setup Distributed Models and Optimizers

Source: https://docs.pytorch.org/docs/2.12/distributed.checkpoint.html

Example setup for FSDP and DDP models with corresponding Adam optimizers.

```python
>>> fsdp_model = FSDP(copy.deepcopy(model))
>>> fsdp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)
>>> ddp_model = DDP(copy.deepcopy(model))
>>> ddp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)
```

--------------------------------

### Install and run TensorBoard

Source: https://docs.pytorch.org/docs/2.12/tensorboard.html

Commands to install the TensorBoard package and launch the server to view logs.

```bash
pip install tensorboard
tensorboard --logdir=runs
```

--------------------------------

### Usage Examples

Source: https://docs.pytorch.org/docs/2.12/_sources/backends.md.txt

Illustrative examples of how to use the `torch.backends.python_native` module.

```APIDOC
## Usage Examples

```python
import torch.backends.python_native as pn

# Query available DSLs
print(pn.available_dsls)  # Example output: ['triton', 'cutedsl']

# Disable all Triton operations
pn.triton.enabled = False

# Temporarily disable CuteDSL operations
with pn.cutedsl.disabled():
    result = model(input)  # CuteDSL ops disabled

# Disable specific operations across all DSLs
pn.disable_operations('scaled_mm', '_flash_attention_forward')

# Query operations for a specific DSL
triton_ops = pn.get_dsl_operations('triton')
```
```

--------------------------------

### Example Input Metadata: pkg.torch.export.graph_signature.InputSpec.kind

Source: https://docs.pytorch.org/docs/2.12/_sources/onnx_export.md.txt

Lists example values for the 'pkg.torch.export.graph_signature.InputSpec.kind' metadata property.

```text
USER_INPUT
```

```text
PARAMETER
```

```text
BUFFER
```

```text
CONSTANT_TENSOR
```

```text
CUSTOM_OBJ
```

```text
TOKEN
```

--------------------------------

### Example Output Metadata: pkg.torch.export.graph_signature.OutputSpec.kind

Source: https://docs.pytorch.org/docs/2.12/_sources/onnx_export.md.txt

Lists example values for the 'pkg.torch.export.graph_signature.OutputSpec.kind' metadata property.

```text
USER_OUTPUT
```

```text
LOSS_OUTPUT
```

```text
BUFFER_MUTATION
```

```text
GRADIENT_TO_PARAMETER
```

```text
GRADIENT_TO_USER_INPUT
```

```text
USER_INPUT_MUTATION
```

```text
TOKEN
```

--------------------------------

### Basic DDP Example with Linear Model

Source: https://docs.pytorch.org/docs/2.12/_sources/notes/ddp.rst.txt

Demonstrates a fundamental Distributed Data Parallel setup using a linear model, including initialization, forward/backward passes, and optimizer steps. Ensure environment variables for master address and port are set.

```python
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
import os
from torch.nn.parallel import DistributedDataParallel as DDP


def example(rank, world_size):
    # create default process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    # create local model
    model = nn.Linear(10, 10).to(rank)
    # construct DDP model
    ddp_model = DDP(model, device_ids=[rank])
    # define loss function and optimizer
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    # forward pass
    outputs = ddp_model(torch.randn(20, 10).to(rank))
    labels = torch.randn(20, 10).to(rank)
    # backward pass
    loss_fn(outputs, labels).backward()
    # update parameters
    optimizer.step()

def main():
    world_size = 2
    mp.spawn(example,
        args=(world_size,),
        nprocs=world_size,
        join=True)

if __name__=="__main__":
    # Environment variables which need to be
    # set when using c10d's default "env"
    # initialization mode.
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29500"
    main()
```

--------------------------------

### Logging output example

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler_troubleshooting.html

Example output showing the result of the reorderable logging configuration.

```bash
$ TORCH_LOGS="graph_breaks" python playground.py
log!
```

--------------------------------

### Example Input Metadata: pkg.torch.export.graph_signature.InputSpec.persistent

Source: https://docs.pytorch.org/docs/2.12/_sources/onnx_export.md.txt

Provides example values for the 'pkg.torch.export.graph_signature.InputSpec.persistent' metadata property.

```text
True
```

```text
False
```

--------------------------------

### RPC Sync Example with Worker 1

Source: https://docs.pytorch.org/docs/2.12/rpc.html

Example of initializing RPC on worker 1 to receive requests.

```python
>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()
```

--------------------------------

### RPC Sync Example with Worker 0

Source: https://docs.pytorch.org/docs/2.12/rpc.html

Example of initializing RPC and performing a synchronous add operation on worker 0.

```python
>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> ret = rpc.rpc_sync("worker1", torch.add, args=(torch.ones(2), 3))
>>> rpc.shutdown()
```

--------------------------------

### Launching Multiple Processes with start_processes

Source: https://docs.pytorch.org/docs/2.12/elastic/multiprocessing.html

Examples demonstrating valid and invalid configurations for launching function and binary entrypoints, including argument handling and environment variable mapping.

```python
log_dir = "/tmp/test"

# ok; two copies of foo: foo("bar0"), foo("bar1")
start_processes(
   name="trainer",
   entrypoint=foo,
   args:{0:("bar0",), 1:("bar1",),
   envs:{0:{}, 1:{}},
   log_dir=log_dir
)

# invalid; envs missing for local rank 1
start_processes(
   name="trainer",
   entrypoint=foo,
   args:{0:("bar0",), 1:("bar1",),
   envs:{0:{}},
   log_dir=log_dir
)

# ok; two copies of /usr/bin/touch: touch file1, touch file2
start_processes(
   name="trainer",
   entrypoint="/usr/bin/touch",
   args:{0:("file1",), 1:("file2",),
   envs:{0:{}, 1:{}},
   log_dir=log_dir
 )

# caution; arguments casted to string, runs:
# echo "1" "2" "3" and echo "[1, 2, 3]"
start_processes(
   name="trainer",
   entrypoint="/usr/bin/echo",
   args:{0:(1,2,3), 1:([1,2,3],),
   envs:{0:{}, 1:{}},
   log_dir=log_dir
 )
```

--------------------------------

### Identify 32-bit Package Installation Failure

Source: https://docs.pytorch.org/docs/2.12/notes/windows.html

Example error output when attempting to install PyTorch on a 32-bit system, which is not supported.

```text
Solving environment: failed

PackagesNotFoundError: The following packages are not available from current channels:

- pytorch

Current channels:
- https://repo.continuum.io/pkgs/main/win-32
- https://repo.continuum.io/pkgs/main/noarch
- https://repo.continuum.io/pkgs/free/win-32
- https://repo.continuum.io/pkgs/free/noarch
- https://repo.continuum.io/pkgs/r/win-32
- https://repo.continuum.io/pkgs/r/noarch
- https://repo.continuum.io/pkgs/pro/win-32
- https://repo.continuum.io/pkgs/pro/noarch
- https://repo.continuum.io/pkgs/msys2/win-32
- https://repo.continuum.io/pkgs/msys2/noarch
```

--------------------------------

### Initialize and use FileStore

Source: https://docs.pytorch.org/docs/2.12/distributed.html

Demonstrates creating a FileStore instance and performing basic key-value operations across distributed processes.

```python
>>> import torch.distributed as dist
>>> store1 = dist.FileStore("/tmp/filestore", 2)
>>> store2 = dist.FileStore("/tmp/filestore", 2)
>>> # Use any of the store methods from either the client or server after initialization
>>> store1.set("first_key", "first_value")
>>> store2.get("first_key")
```

--------------------------------

### Example of Guard Installation with Constant Input

Source: https://docs.pytorch.org/docs/2.12/_sources/user_guide/torch_compiler/torch.compiler_dynamo_deepdive.md.txt

This snippet demonstrates how a constant string input installs guards for type and value equality when using `@torch.compile`. Run with `TORCH_LOGS=guards` to see the generated guards.

```python
import torch

@torch.compile
def fn(a, b):
    return a * len(b)

fn(torch.arange(10), "Hello")
```

--------------------------------

### Example ConstraintViolation error message

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/export/api_reference.html

Example of suggested fixes provided in a ConstraintViolation error message during export.

```text
Suggested fixes:

    dim = Dim('dim', min=3, max=6)  # this just refines the dim's range
    dim = 4  # this specializes to a constant
    dy = dx + 1  # dy was specified as an independent dim, but is actually tied to dx with this relation
```

--------------------------------

### Initialize FSDP with custom parameter initialization

Source: https://docs.pytorch.org/docs/2.12/fsdp.html

Demonstrates how to use param_init_fn to initialize modules on the meta device and how to integrate with torchdistX for deferred initialization.

```python
>>> module = MyModule(device="meta")
>>> def my_init_fn(module: nn.Module):
>>>     # E.g. initialize depending on the module type
>>>     ...
>>> fsdp_model = FSDP(module, param_init_fn=my_init_fn, auto_wrap_policy=size_based_auto_wrap_policy)
>>> print(next(fsdp_model.parameters()).device) # current CUDA device
>>> # With torchdistX
>>> module = deferred_init.deferred_init(MyModule, device="cuda")
>>> # Will initialize via deferred_init.materialize_module().
>>> fsdp_model = FSDP(module, auto_wrap_policy=size_based_auto_wrap_policy)
```

--------------------------------

### Constrain as Value Example

Source: https://docs.pytorch.org/docs/2.12/generated/exportdb/torch.dynamic-value.html

Demonstrates using torch._check to constrain a dynamic value to guide control flow during export.

```python
# mypy: allow-untyped-defs
import torch


class ConstrainAsValueExample(torch.nn.Module):
    """
    If the value is not known at tracing time, you can provide hint so that we
    can trace further. Please look at torch._check API.
    """

    def forward(self, x, y):
        a = x.item()
        torch._check(a >= 0)
        torch._check(a <= 5)

        if a < 6:
            return y.sin()
        return y.cos()


example_args = (torch.tensor(4), torch.randn(5, 5))
tags = {
    "torch.dynamic-value",
    "torch.escape-hatch",
}
model = ConstrainAsValueExample()


torch.export.export(model, example_args)
```

```text
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "i64[]", y: "f32[5, 5]"):
                item: "Sym(u0)" = torch.ops.aten.item.default(x);  x = None
            ge_1: "Sym(u0 >= 0)" = item >= 0
            _assert_scalar_default = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u0 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_default = None
            le_1: "Sym(u0 <= 5)" = item <= 5;  item = None
            _assert_scalar_default_1 = torch.ops.aten._assert_scalar.default(le_1, "Runtime assertion failed for expression u0 <= 5 on node 'le_1'");  le_1 = _assert_scalar_default_1 = None

                sin: "f32[5, 5]" = torch.ops.aten.sin.default(y);  y = None
            return (sin,)

Graph signature:
    # inputs
    x: USER_INPUT
    y: USER_INPUT

    # outputs
    sin: USER_OUTPUT

Range constraints: {u0: VR[0, 5], u1: VR[0, 5]}
```

--------------------------------

### Generating Logarithmically Spaced Tensors

Source: https://docs.pytorch.org/docs/2.12/generated/torch.logspace.html

Examples demonstrating how to use torch.logspace with various start, end, steps, and base parameters.

```python
>>> torch.logspace(start=-10, end=10, steps=5)
tensor([ 1.0000e-10,  1.0000e-05,  1.0000e+00,  1.0000e+05,  1.0000e+10])
>>> torch.logspace(start=0.1, end=1.0, steps=5)
tensor([  1.2589,   2.1135,   3.5481,   5.9566,  10.0000])
>>> torch.logspace(start=0.1, end=1.0, steps=1)
tensor([1.2589])
>>> torch.logspace(start=2, end=2, steps=1, base=2)
tensor([4.0])
```

--------------------------------

### Generating evenly spaced tensors with torch.linspace

Source: https://docs.pytorch.org/docs/2.12/generated/torch.linspace.html

Examples demonstrating how to create tensors with a specific number of steps between start and end values.

```python
>>> torch.linspace(3, 10, steps=5)
tensor([  3.0000,   4.7500,   6.5000,   8.2500,  10.0000])
>>> torch.linspace(-10, 10, steps=5)
tensor([-10.,  -5.,   0.,   5.,  10.])
>>> torch.linspace(start=-10, end=10, steps=5)
tensor([-10.,  -5.,   0.,   5.,  10.])
>>> torch.linspace(start=-10, end=10, steps=1)
tensor([-10.])
```

--------------------------------

### KLDivLoss usage examples

Source: https://docs.pytorch.org/docs/2.12/generated/torch.nn.KLDivLoss.html

Demonstrates initializing KLDivLoss with different reduction and log_target configurations.

```python
>>> kl_loss = nn.KLDivLoss(reduction="batchmean")
>>> # input should be a distribution in the log space
>>> input = F.log_softmax(torch.randn(3, 5, requires_grad=True), dim=1)
>>> # Sample a batch of distributions. Usually this would come from the dataset
>>> target = F.softmax(torch.rand(3, 5), dim=1)
>>> output = kl_loss(input, target)
>>>
>>> kl_loss = nn.KLDivLoss(reduction="batchmean", log_target=True)
>>> log_target = F.log_softmax(torch.rand(3, 5), dim=1)
>>> output = kl_loss(input, log_target)
```

--------------------------------

### Initialize and draw from SobolEngine

Source: https://docs.pytorch.org/docs/2.12/generated/torch.quasirandom.SobolEngine.html

Demonstrates initializing a SobolEngine with a specified dimension and drawing a sequence of points.

```python
>>> soboleng = torch.quasirandom.SobolEngine(dimension=5)
>>> soboleng.draw(3)
tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
        [0.7500, 0.2500, 0.2500, 0.2500, 0.7500]])
```

--------------------------------

### Accessing the real part of a complex tensor

Source: https://docs.pytorch.org/docs/2.12/_sources/generated/torch.Tensor.real.rst.txt

This example demonstrates how to get the real part of a complex tensor using the `.real` attribute.

```APIDOC
## Accessing the real part of a complex tensor

### Description

This snippet shows how to retrieve the real component of a complex-valued tensor.

### Method

Attribute access

### Endpoint

`Tensor.real`

### Parameters

This attribute does not take any parameters.

### Request Example

```python
import torch

# Create a complex tensor
complex_tensor = torch.tensor([1+2j, 3+4j, 5+6j])

# Get the real part
real_part = complex_tensor.real

print(real_part)
```

### Response

#### Success Response

Returns a tensor containing the real parts of the input tensor.

#### Response Example

```
tensor([1., 3., 5.])
```
```

--------------------------------

### Prepare a model for quantization using FX

Source: https://docs.pytorch.org/docs/2.12/generated/torch.ao.quantization.quantize_fx.prepare_fx.html

Demonstrates the end-to-end workflow of defining a model, configuring quantization settings, and using prepare_fx to insert observers.

```python
import torch
from torch.ao.quantization import get_default_qconfig_mapping
from torch.ao.quantization.quantize_fx import prepare_fx

class Submodule(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = torch.nn.Linear(5, 5)
    def forward(self, x):
        x = self.linear(x)
        return x

class M(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = torch.nn.Linear(5, 5)
        self.sub = Submodule()

    def forward(self, x):
        x = self.linear(x)
        x = self.sub(x) + x
        return x

# initialize a floating point model
float_model = M().eval()

# define calibration function
def calibrate(model, data_loader):
    model.eval()
    with torch.no_grad():
        for image, target in data_loader:
            model(image)

# qconfig is the configuration for how we insert observers for a particular
# operator
# qconfig = get_default_qconfig("fbgemm")
# Example of customizing qconfig:
# qconfig = torch.ao.quantization.QConfig(
#    activation=MinMaxObserver.with_args(dtype=torch.qint8),
#    weight=MinMaxObserver.with_args(dtype=torch.qint8))
# `activation` and `weight` are constructors of observer module

# qconfig_mapping is a collection of quantization configurations, user can
# set the qconfig for each operator (torch op calls, functional calls, module calls)
# in the model through qconfig_mapping
# the following call will get the qconfig_mapping that works best for models
# that target "fbgemm" backend
qconfig_mapping = get_default_qconfig_mapping("fbgemm")

# We can customize qconfig_mapping in different ways.
# e.g. set the global qconfig, which means we will use the same qconfig for
# all operators in the model, this can be overwritten by other settings
# qconfig_mapping = QConfigMapping().set_global(qconfig)
# e.g. quantize the linear submodule with a specific qconfig
# qconfig_mapping = QConfigMapping().set_module_name("linear", qconfig)
# e.g. quantize all nn.Linear modules with a specific qconfig
# qconfig_mapping = QConfigMapping().set_object_type(torch.nn.Linear, qconfig)
# for a more complete list, please see the docstring for :class:`torch.ao.quantization.QConfigMapping`
# argument

# example_inputs is a tuple of inputs, that is used to infer the type of the
# outputs in the model
# currently it's not used, but please make sure model(*example_inputs) runs
example_inputs = (torch.randn(1, 3, 224, 224),)

# TODO: add backend_config after we split the backend_config for fbgemm and qnnpack
# e.g. backend_config = get_default_backend_config("fbgemm")
# `prepare_fx` inserts observers in the model based on qconfig_mapping and
# backend_config. If the configuration for an operator in qconfig_mapping
# is supported in the backend_config (meaning it's supported by the target
# hardware), we'll insert observer modules according to the qconfig_mapping
# otherwise the configuration in qconfig_mapping will be ignored
#
# Example:
# in qconfig_mapping, user sets linear module to be quantized with quint8 for
# activation and qint8 for weight:
# qconfig = torch.ao.quantization.QConfig(
#     observer=MinMaxObserver.with_args(dtype=torch.quint8),
#     weight=MinMaxObserver.with-args(dtype=torch.qint8))
# Note: current qconfig api does not support setting output observer, but
# we may extend this to support these more fine grained control in the
# future

# qconfig_mapping = QConfigMapping().set_object_type(torch.nn.Linear, qconfig)
# in backend config, linear module also supports in this configuration:
# weighted_int8_dtype_config = DTypeConfig(
#   input_dtype=torch.quint8,
#   output_dtype=torch.quint8,
#   weight_dtype=torch.qint8,
#   bias_type=torch.float)

# linear_pattern_config = BackendPatternConfig(torch.nn.Linear) \
#    .set_observation_type(ObservationType.OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT) \
#    .add_dtype_config(weighted_int8_dtype_config) \
#    ...

# backend_config = BackendConfig().set_backend_pattern_config(linear_pattern_config)
# `prepare_fx` will check that the setting requested by suer in qconfig_mapping
# is supported by the backend_config and insert observers and fake quant modules
# in the model
prepared_model = prepare_fx(float_model, qconfig_mapping, example_inputs)
# Run calibration
calibrate(prepared_model, sample_inference_data)
```

--------------------------------

### Higher-Precision Reduction Example Setup

Source: https://docs.pytorch.org/docs/2.12/_sources/symmetric_memory.md.txt

This snippet shows the necessary imports for utilizing higher-precision reduction with symmetric memory tensors in PyTorch distributed operations.

```python
import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem
```

--------------------------------

### Create a C++ extension with setuptools

Source: https://docs.pytorch.org/docs/2.12/cpp_extension.html

Use CppExtension to define a C++ extension module within a setup.py file.

```python
>>> from setuptools import setup
>>> from torch.utils.cpp_extension import BuildExtension, CppExtension
>>> setup(
...     name='extension',
...     ext_modules=[
...         CppExtension(
...             name='extension',
...             sources=['extension.cpp'],
...             extra_compile_args=['-g'],
...             extra_link_args=['-Wl,--no-as-needed', '-lm'])
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })
```

--------------------------------

### DDP Initialization Logs

Source: https://docs.pytorch.org/docs/2.12/_sources/distributed.md.txt

Example of debug logs rendered at initialization time when TORCH_DISTRIBUTED_DEBUG is set to INFO or DETAIL. These logs provide details about the DDP setup.

```log
I0607 16:10:35.739390 515217 logger.cpp:173] [Rank 0]: DDP Initialized with:
broadcast_buffers: 1
bucket_cap_bytes: 26214400
find_unused_parameters: 0
gradient_as_bucket_view: 0
is_multi_device_module: 0
iteration: 0
num_parameter_tensors: 2
output_device: 0
rank: 0
total_parameter_size_bytes: 440
backend_name: nccl
bucket_sizes: 440
cuda_visible_devices: N/A
device_ids: 0
dtypes: float
master_addr: localhost
master_port: 29501
module_name: TwoLinLayerNet
nccl_async_error_handling: N/A
nccl_blocking_wait: N/A
nccl_debug: WARN
nccl_ib_timeout: N/A
nccl_nthreads: N/A
nccl_socket_ifname: N/A
torch_distributed_debug: INFO

```

--------------------------------

### Example Python script for tracing

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler_troubleshooting.html

A simple function decorated with torch.compile to demonstrate tracing behavior.

```python
import torch

def g(x, y):
    return x + y

@torch.compile(backend="eager")
def f(x):
    x = torch.sin(x)
    x = g(x, x)
    return x

f(torch.ones(3, 3))
```

--------------------------------

### Initialize Modules with Custom Settings

Source: https://docs.pytorch.org/docs/2.12/notes/modules.html

Demonstrates initializing modules directly on GPU, using 16-bit floating point, and skipping default initialization for custom schemes.

```python
# Initialize module directly onto GPU.
m = nn.Linear(5, 3, device='cuda')

# Initialize module with 16-bit floating point parameters.
m = nn.Linear(5, 3, dtype=torch.half)

# Skip default parameter initialization and perform custom (e.g. orthogonal) initialization.
m = torch.nn.utils.skip_init(nn.Linear, 5, 3)
nn.init.orthogonal_(m.weight)
```

--------------------------------

### RPC Sync Example for Distributed Model

Source: https://docs.pytorch.org/docs/2.12/_sources/rpc/distributed_autograd.rst.txt

Demonstrates a simple distributed model setup using RPC, where a function is called synchronously on a remote worker. Ensure RPC is initialized before use.

```python
import torch
import torch.distributed.rpc as rpc

def my_add(t1, t2):
  return torch.add(t1, t2)

# On worker 0:
t1 = torch.rand((3, 3), requires_grad=True)
t2 = torch.rand((3, 3), requires_grad=True)

# Perform some computation remotely.
t3 = rpc.rpc_sync("worker1", my_add, args=(t1, t2))

# Perform some computation locally based on remote result.
t4 = torch.rand((3, 3), requires_grad=True)
t5 = torch.mul(t3, t4)

# Compute some loss.
loss = t5.sum()
```

--------------------------------

### Implementing BUILD_LIST Bytecode

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler_dynamo_deepdive.html

Example implementation of the BUILD_LIST bytecode within InstructionTranslatorBase.

```python
def BUILD_LIST(self, inst):
    items = self.popn(inst.argval)
    self.push(ListVariable(items, mutation_type=ValueMutationNew()))
```

--------------------------------

### Applying torch.compile to Models

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler_troubleshooting.html

Examples of applying torch.compile for inference, training, and distributed training scenarios.

```python
# inference
model = ...
opt_model = torch.compile(model)

for _ in range(N_ITERS):
    inp = ...
    out = opt_model(inp)
```

```python
# training
model = ...
opt = torch.optim.Adam(model.parameters())

@torch.compile
def train(mod, data):
    opt.zero_grad(True)
    pred = mod(data[0])
    loss = torch.nn.CrossEntropyLoss()(pred, data[1])
    loss.backward()
    opt.step()

for _ in range(N_ITERS):
    inp = ...
    train(model, inp)
```

```python
# DistributedDataParallel
model = ...
opt_model = torch.compile(model)
model_ddp = DistributedDataParallel(opt_model, ...)

for _ in range(N_ITERS):
    inp = ...
    out = model_ddp(inp)
```

--------------------------------

### Inspecting Generated Guards

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler_troubleshooting.html

Example showing how to view the runtime guards generated by torch.compile for a specific function.

```python
import torch

@torch.compile
def fn(x):
    return x + 1

fn(torch.ones(3, 3))
```

```bash
$ TORCH_LOGS="guards" python playground.py
GUARDS:

TREE_GUARD_MANAGER:
+- RootGuardManager
| +- DEFAULT_DEVICE: utils_device.CURRENT_DEVICE == None                           # _dynamo/output_graph.py:471 in init_ambient_guards
| +- GLOBAL_STATE: ___check_global_state()
| +- TORCH_FUNCTION_MODE_STACK: ___check_torch_function_mode_stack()
| +- GuardManager: source=L['x'], accessed_by=DictGetItemGuardAccessor(x)
| | +- TENSOR_MATCH: check_tensor(L['x'], Tensor, DispatchKeySet(CPU, BackendSelect, ADInplaceOrView, AutogradCPU), torch.float32, device=None, requires_grad=False, size=[3, 3], stride=[3, 1])  # return x + 1  # playground.py:6 in fn
| | +- NO_HASATTR: hasattr(L['x'], '_dynamo_dynamic_indices') == False           # return x + 1  # playground.py:6 in fn
```

--------------------------------

### setup_context(ctx, inputs, output)

Source: https://docs.pytorch.org/docs/2.12/generated/torch.autograd.function.NestedIOFunction.html

Sets up the context for the backward pass when the forward pass signature does not include the context object.

```APIDOC
### Method
static setup_context(ctx, inputs, output)

### Description
Used when overriding forward without the ctx argument. This method is responsible for saving tensors or setting context attributes required for the backward pass.
```

--------------------------------

### Training with FP32 on Intel GPU

Source: https://docs.pytorch.org/docs/2.12/_sources/notes/get_start_xpu.rst.txt

Example of a training workflow using the CIFAR10 dataset and a ResNet50 model on an Intel GPU with FP32 precision. It includes data loading, model setup, loss calculation, and optimization steps.

```python
import torch
import torchvision

LR = 0.001
DOWNLOAD = True
DATA = "datasets/cifar10/"

transform = torchvision.transforms.Compose(
    [
        torchvision.transforms.Resize((224, 224)),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)
train_dataset = torchvision.datasets.CIFAR10(
    root=DATA,
    train=True,
    transform=transform,
    download=DOWNLOAD,
)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=128)
train_len = len(train_loader)

model = torchvision.models.resnet50()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=LR, momentum=0.9)
model.train()
model = model.to("xpu")
criterion = criterion.to("xpu")

print(f"Initiating training")
for batch_idx, (data, target) in enumerate(train_loader):
    data = data.to("xpu")
    target = target.to("xpu")
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    if (batch_idx + 1) % 10 == 0:
        iteration_loss = loss.item()
        print(f"Iteration [{batch_idx+1}/{train_len}], Loss: {iteration_loss:.4f}")
torch.save(
    {

```

--------------------------------

### Defining Sources for Variables

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler_dynamo_deepdive.html

Example illustrating how local variables and intermediate variables are tracked as sources.

```python
def foo(x: Tensor, y: List[Tensor]):
    a = x * y[0]
    return a * x
```

--------------------------------

### Initialize Model and State Dictionary

Source: https://docs.pytorch.org/docs/2.12/distributed.checkpoint.html

Setup a model instance and a state dictionary for checkpointing.

```python
>>> my_model = MyModule()
```

```python
>>> state_dict = {"model": my_model}
```

--------------------------------

### Inspect Dynamo state via Pdb

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler_troubleshooting.html

Example output and commands available when a comptime breakpoint is triggered during execution.

```bash
$ python playground.py
--Return--
> /data/users/williamwen/pytorch/torch/_dynamo/comptime.py(392)inner()->None
-> builtins.breakpoint()
(Pdb) ctx.print_bt()
File "/data/users/williamwen/pytorch/playground.py", line 7, in f
    comptime.breakpoint()

(Pdb) ctx.print_locals()
x = FakeTensor(..., size=(3, 3))
y = FakeTensor(..., size=(3, 3))
(Pdb) bt
...
/data/users/williamwen/pytorch/torch/_dynamo/symbolic_convert.py(826)call_function()
-> self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
/data/users/williamwen/pytorch/torch/_dynamo/variables/misc.py(331)call_function()
-> func(ComptimeContext(tx))
> /data/users/williamwen/pytorch/torch/_dynamo/comptime.py(392)inner()->None
-> builtins.breakpoint()
(Pdb) ctx.print_graph()



def forward(self, L_x_: "f32[3, 3]"):
    l_x_ = L_x_

    # File: /data/users/williamwen/pytorch/playground.py:6 in f, code: y = x + 1
    y: "f32[3, 3]" = l_x_ + 1;  l_x_ = y = None
```

--------------------------------

### Initialize SummaryWriter with different configurations

Source: https://docs.pytorch.org/docs/2.12/tensorboard.html

Demonstrates various ways to initialize the SummaryWriter, including default, custom directory, and comment-based naming.

```python
from torch.utils.tensorboard import SummaryWriter

# create a summary writer with automatically generated folder name.
writer = SummaryWriter()
# folder location: runs/May04_22-14-54_s-MacBook-Pro.local/

# create a summary writer using the specified folder name.
writer = SummaryWriter("my_experiment")
# folder location: my_experiment

# create a summary writer with comment appended.
writer = SummaryWriter(comment="LR_0.1_BATCH_16")
# folder location: runs/May04_22-14-54_s-MacBook-Pro.localLR_0.1_BATCH_16/
```

--------------------------------

### Saving and Reloading PowerSGD Hook State

Source: https://docs.pytorch.org/docs/2.12/_sources/ddp_comm_hooks.md.txt

This example demonstrates how to save the state of a DDP model with a registered PowerSGD communication hook and then reload it. It covers model definition, distributed setup, hook registration, checkpointing, and restoring the model and hook state.

```python
import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel
from torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook as powerSGD

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(24,24)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(24,12)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def run_demo(demo_fn, world_size):
    mp.spawn(
        demo_fn,
        args=(world_size,),
        nprocs=world_size,
        join=True)

def demo_serialization(rank, world_size):
    setup(rank, world_size)

    CHECKPOINT = tempfile.gettempdir() + "/checkpoint.pt"

    model = SimpleModel().to(rank)
    ddp_model = DistributedDataParallel(model, device_ids=[rank])

    powersgd_hook = powerSGD.powerSGD_hook
    powersgd_state = powerSGD.PowerSGDState(process_group=None)

    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    ddp_model.register_comm_hook(powersgd_state, powersgd_hook)

    state = {
        'state_dict': ddp_model.state_dict(),
        'comm_hook': powersgd_hook,
        'comm_hook_state': powersgd_state}

    if rank == 0:
        torch.save(state, CHECKPOINT)

    dist.barrier()
    map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
    checkpoint = torch.load(CHECKPOINT, map_location=map_location)

    new_ddp_model = DistributedDataParallel(SimpleModel().to(rank), device_ids=[rank])
    new_ddp_model.load_state_dict(checkpoint['state_dict'])
    powersgd_hook = checkpoint['comm_hook']
    powersgd_state = checkpoint['comm_hook_state']

    new_ddp_model.register_comm_hook(powersgd_state, powersgd_hook)

    if rank == 0:
        os.remove(CHECKPOINT)

    cleanup()

if __name__ == "__main__":
    n_gpus = torch.cuda.device_count()
    assert n_gpus >= 2, f"Requires at least 2 GPUs to run, but got {n_gpus}"
    world_size = n_gpus
    run_demo(demo_serialization, world_size)

```

--------------------------------

### Generating Guards for GetItemSource

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler_dynamo_deepdive.html

Example showing how GetItemSource generates guards for list elements.

```python
import torch

@torch.compile
def fn(x, l):
    return x * len(l[0])

fn(torch.randn(8), ["Hi", "Hello"])
```

--------------------------------

### torch.export.export

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/export/api_reference.html

Traces an nn.Module with example inputs to produce an ExportedProgram, handling static and dynamic shape constraints.

```APIDOC
## torch.export.export

### Description
`export()` takes any nn.Module along with example inputs, and produces a traced graph representing only the Tensor computation of the function in an Ahead-of-Time (AOT) fashion. It eliminates Python control flow and records shape constraints.

### Parameters
- **mod** (nn.Module) - The module to be exported.
- **args** (tuple) - Example inputs for the module.
- **kwargs** (dict, optional) - Keyword arguments for the module.
- **dynamic_shapes** (dict, optional) - Specification of dynamic dimensions using the Dim() API.
- **strict** (bool, optional) - Whether to enforce strict tracing.
- **preserve_module_call_signature** (tuple, optional) - Signatures to preserve during export.
- **prefer_deferred_runtime_asserts_over_guards** (bool, optional) - Preference for runtime assertions over guards.
```

--------------------------------

### example_arguments()

Source: https://docs.pytorch.org/docs/2.12/onnx_export.html

Returns example arguments for the model's forward method, used for tracing and ONNX export.

```APIDOC
## example_arguments()

### Description
Returns a tuple of positional arguments and a dictionary of keyword arguments to be used as sample inputs for the model's forward method during tracing or ONNX export.

### Returns
- **tuple** - Positional arguments for the forward method.
- **dict | None** - Keyword arguments for the forward method.
```

--------------------------------

### Train with torch.compile on XPU

Source: https://docs.pytorch.org/docs/2.12/_sources/notes/get_start_xpu.rst.txt

This example shows how to train a ResNet50 model on CIFAR-10 using torch.compile for graph optimization on an XPU device. It includes data loading, model setup, training loop, and checkpoint saving. The model is compiled before the training loop begins.

```python
import torch
import torchvision

LR = 0.001
DOWNLOAD = True
DATA = "datasets/cifar10/"

transform = torchvision.transforms.Compose(
    [
        torchvision.transforms.Resize((224, 224)),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)
train_dataset = torchvision.datasets.CIFAR10(
    root=DATA,
    train=True,
    transform=transform,
    download=DOWNLOAD,
)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=128)
train_len = len(train_loader)

model = torchvision.models.resnet50()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=LR, momentum=0.9)
model.train()
model = model.to("xpu")
criterion = criterion.to("xpu")
model = torch.compile(model)

print(f"Initiating training with torch compile")
for batch_idx, (data, target) in enumerate(train_loader):
    data = data.to("xpu")
    target = target.to("xpu")
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    if (batch_idx + 1) % 10 == 0:
        iteration_loss = loss.item()
        print(f"Iteration [{batch_idx+1}/{train_len}], Loss: {iteration_loss:.4f}")
torch.save(
    {
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
    },
    "checkpoint.pth",
)

print("Execution finished")
```

--------------------------------

### GetItemSource Guard Output

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler_dynamo_deepdive.html

The resulting guards generated by the GetItemSource example.

```text
___check_type_id(L['l'], 94439025877664)
len(L['l']) == 2
___check_type_id(L['l'][0], 94439025840192)
L['l'][0] == 'Hi'
___check_type_id(L['l'][1], 94439025840192)
L['l'][1] == 'Hello'
```

--------------------------------

### View launch utility help

Source: https://docs.pytorch.org/docs/2.12/distributed.html

Display available command-line arguments for the launch module.

```bash
python -m torch.distributed.launch --help
```

--------------------------------

### Exponential Moving Average (EMA) Example

Source: https://docs.pytorch.org/docs/2.12/_sources/optim.md.txt

Illustrates training a model using EMA with a decay rate of 0.999. The EMA model starts collecting averages immediately. The training loop updates model parameters and then updates the EMA model. Finally, it updates BN statistics and uses the EMA model for predictions.

```python
>>> loader, optimizer, model, loss_fn = ...
>>> ema_model = torch.optim.swa_utils.AveragedModel(model, \
>>>             multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(0.999))
>>>
>>> for epoch in range(300):
>>>       for input, target in loader:
>>>           optimizer.zero_grad()
>>>           loss_fn(model(input), target).backward()
>>>           optimizer.step()
>>>           ema_model.update_parameters(model)
>>>
>>> # Update bn statistics for the ema_model at the end
>>> torch.optim.swa_utils.update_bn(loader, ema_model)
>>> # Use ema_model to make predictions on test data
>>> preds = ema_model(test_input)
```

--------------------------------

### Recompilation in LR Schedulers

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler_troubleshooting.html

Example showing how LR scheduler updates can trigger recompilations due to constant guarding of the learning rate.

```python
import torch

mod = torch.nn.Linear(3, 3)
opt = torch.optim.Adam(mod.parameters(), lr=0.01)
sched = torch.optim.lr_scheduler.ExponentialLR(opt, 0.9)

@torch.compile
def fn(inp):
    opt.zero_grad(True)
    out = mod(inp).sum()
    out.backward()
    opt.step()
    sched.step()

for i in range(1, 10):
    fn(torch.ones(3, 3))

```

```bash
$ TORCH_LOGS="recompiles" python playground.py
Recompiling function step in /data/users/williamwen/pytorch/torch/optim/adam.py:189
    triggered by the following guard failure(s):
    - 3/7: L['self'].param_groups[0]['lr'] == 0.004782969000000002
    - 3/6: L['self'].param_groups[0]['lr'] == 0.005314410000000002
    - 3/5: L['self'].param_groups[0]['lr'] == 0.005904900000000002
    - 3/4: L['self'].param_groups[0]['lr'] == 0.006561000000000002
    - 3/3: L['self'].param_groups[0]['lr'] == 0.007290000000000001
    - 3/2: L['self'].param_groups[0]['lr'] == 0.008100000000000001
    - 3/1: L['self'].param_groups[0]['lr'] == 0.009000000000000001
    - 3/0: L['self'].param_groups[0]['lr'] == 0.01
torch._dynamo hit config.recompile_limit (8)
    function: 'step' (/data/users/williamwen/pytorch/torch/optim/adam.py:189)
    last reason: 3/0: L['self'].param_groups[0]['lr'] == 0.01

```

--------------------------------

### torch.Tensor.narrow

Source: https://docs.pytorch.org/docs/2.12/_sources/generated/torch.Tensor.narrow.rst.txt

Returns a tensor that is a narrowed version of the input tensor.

This operation returns a view on the original tensor. No memory is copied.

Args:
    dimension (int): The dimension along which to narrow.
    start (int): The starting index of the slice.
    length (int): The length of the slice.

Example:

.. code-block:: python

    >>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
    >>> x.narrow(0, 1, 2)
    tensor([[4, 5, 6],
            [7, 8, 9]])
    >>> x.narrow(1, 0, 1)
    tensor([[1],
            [4],
            [7]])

```APIDOC
## torch.Tensor.narrow

### Description
Returns a tensor that is a narrowed version of the input tensor. This operation returns a view on the original tensor. No memory is copied.

### Parameters
* **dimension** (int) - The dimension along which to narrow.
* **start** (int) - The starting index of the slice.
* **length** (int) - The length of the slice.

### Example
```python
x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print(x.narrow(0, 1, 2))
# tensor([[4, 5, 6],
#         [7, 8, 9]])
print(x.narrow(1, 0, 1))
# tensor([[1],
#         [4],
#         [7]])
```
```

--------------------------------

### Initialize and use OneCycleLR

Source: https://docs.pytorch.org/docs/2.12/generated/torch.optim.lr_scheduler.OneCycleLR.html

Demonstrates setting up the scheduler with an optimizer and calling step() after each batch.

```python
>>> data_loader = torch.utils.data.DataLoader(...)
>>> optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
>>> scheduler = torch.optim.lr_scheduler.OneCycleLR(
...     optimizer, max_lr=0.01, steps_per_epoch=len(data_loader), epochs=10
... )
>>> for epoch in range(10):
>>>     for batch in data_loader:
>>>         train_batch(...)
>>>         optimizer.step()
>>>         scheduler.step()
```

--------------------------------

### Diagnosing Recompilation

Source: https://docs.pytorch.org/docs/2.12/user_guide/torch_compiler/torch.compiler_troubleshooting.html

Example demonstrating how to identify when a function is recompiled due to guard failures, such as tensor shape mismatches.

```python
import torch

@torch.compile
def fn(x):
    return x + 1

fn(torch.ones(3, 3))
fn(torch.ones(4, 4))
```

```bash
$ TORCH_LOGS="recompiles" python playground.py
Recompiling function fn in /data/users/williamwen/pytorch/playground.py:3
    triggered by the following guard failure(s):
    - 0/0: tensor 'L['x']' size mismatch at index 0. expected 3, actual 4
```
Tokens:
Raw