PyTorch

https://docs.pytorch.org/docs/2.13
Admin
PyTorch is an open-source machine learning framework and optimized tensor library for deep learning...
Tokens:1,531,716
Snippets:15,060
Trust Score:10
Update:1 month ago
Show doc for...
Tokens:
Raw
### Install and run TensorBoard

Source: https://docs.pytorch.org/docs/2.13/_sources/tensorboard.md.txt

Commands to install the TensorBoard package and launch the server to visualize logged data.

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

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

### Matrix Inversion Examples

Source: https://docs.pytorch.org/docs/2.13/generated/torch.linalg.inv.html

Examples demonstrating the inversion of standard, batched, and complex matrices using torch.linalg.inv.

```python
>>> A = torch.randn(4, 4)
>>> Ainv = torch.linalg.inv(A)
>>> torch.dist(A @ Ainv, torch.eye(4))
tensor(1.1921e-07)

>>> A = torch.randn(2, 3, 4, 4)  # Batch of matrices
>>> Ainv = torch.linalg.inv(A)
>>> torch.dist(A @ Ainv, torch.eye(4))
tensor(1.9073e-06)

>>> A = torch.randn(4, 4, dtype=torch.complex128)  # Complex matrix
>>> Ainv = torch.linalg.inv(A)
>>> torch.dist(A @ Ainv, torch.eye(4))
tensor(7.5107e-16, dtype=torch.float64)
```

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

### Example Package Output

Source: https://docs.pytorch.org/docs/2.13/package.html

The resulting file structure and output after running the packaging script.

```text
─── foo_package
    ├── foo-generated
    │   ├── _0
    │   │   └── foo.txt
    │   └── _1
    │       └── foo.txt
    ├── foo_collection
    │   ├── foo1.pkl
    │   └── foo2.pkl
    └── foo.py

foo_1 string: 'foo_1 initial string, with reduction modification!'
foo_1 export time: 9857706.650140837
foo_1 import time: 9857706.652698385
```

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

### Launching Multiple Processes with start_processes

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

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

```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
 )
```

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

### Define a function for compilation

Source: https://docs.pytorch.org/docs/2.13/_sources/user_guide/torch_compiler/torch.compiler_troubleshooting.md.txt

Example function setup 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 and Use TCPStore

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

Demonstrates setting up a TCPStore server and client to share key-value pairs across processes.

```python
>>> import torch.distributed as dist
>>> from datetime import timedelta
>>> # Run on process 1 (server)
>>> server_store = dist.TCPStore("127.0.0.1", 1234, 2, True, timedelta(seconds=30))
>>> # Run on process 2 (client)
>>> client_store = dist.TCPStore("127.0.0.1", 1234, 2, False)
>>> # Use any of the store methods from either the client or server after initialization
>>> server_store.set("first_key", "first_value")
>>> client_store.get("first_key")
```

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

### Generating Tensors with torch.arange

Source: https://docs.pytorch.org/docs/2.13/generated/torch.arange.html

Examples demonstrating the creation of tensors using different start, end, and step parameters.

```python
>>> torch.arange(5)
tensor([ 0,  1,  2,  3,  4])
>>> torch.arange(1, 4)
tensor([ 1,  2,  3])
>>> torch.arange(1, 2.5, 0.5)
tensor([ 1.0000,  1.5000,  2.0000])
```

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

### Generating logarithmic tensors with torch.logspace

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

Examples demonstrating the creation of tensors with different 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])
```

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

### Initialize and use RemoteModule across processes

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

Example setup for running RemoteModule across two workers, demonstrating asynchronous invocation.

```python
>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> from torch import nn, Tensor
>>> from torch.distributed.nn.api.remote_module import RemoteModule
>>>
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> remote_linear_module = RemoteModule(
>>>     "worker1/cpu", nn.Linear, args=(20, 30),
>>> )
>>> input = torch.randn(128, 20)
>>> ret_fut = remote_linear_module.forward_async(input)
>>> ret = ret_fut.wait()
>>> rpc.shutdown()
```

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

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

### setup_context(ctx, inputs, output)

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

Defines the context setup for the backward pass when using the alternative forward signature.

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

### Description
Sets up the context for the backward pass. This is used when overriding forward with the signature forward(*args, **kwargs).

### Parameters
- **ctx** (Context) - Required - The context object.
- **inputs** (Any) - Required - The inputs to the forward function.
- **output** (Any) - Required - The output of the forward function.
```

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

### Initialize and Use FileStore

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

Demonstrates using a file-based store to share data between 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")
```

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

### Compute local autograd dependencies

Source: https://docs.pytorch.org/docs/2.13/_sources/rpc/distributed_autograd.md.txt

Example of a standard local autograd graph setup where the engine calculates dependencies based on node usage.

```python
import torch
a = torch.rand((3, 3), requires_grad=True)
b = torch.rand((3, 3), requires_grad=True)
c = torch.rand((3, 3), requires_grad=True)
d = a + b
e = b * c
d.sum().backward()
```

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

### Initialize Tensors and Networks on Device

Source: https://docs.pytorch.org/docs/2.13/_sources/notes/cuda.md.txt

Demonstrates creating tensors and moving networks to the device determined by the configuration.

```python
x = torch.empty((8, 42), device=args.device)
net = Network().to(device=args.device)
```

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

### FakeTensor instance example

Source: https://docs.pytorch.org/docs/2.13/_sources/user_guide/torch_compiler/export/ir_spec.md.txt

An example instantiation of a FakeTensor object.

```python
FakeTensor(dtype=torch.int, size=[2,], device=CPU)
```

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

### Compute rfftfreq examples

Source: https://docs.pytorch.org/docs/2.13/generated/torch.fft.rfftfreq.html

Examples showing the output of rfftfreq for different input lengths.

```python
>>> torch.fft.rfftfreq(5)
tensor([0.0000, 0.2000, 0.4000])
```

```python
>>> torch.fft.rfftfreq(4)
tensor([0.0000, 0.2500, 0.5000])
```

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

### Prepare a model for QAT using FX Graph Mode

Source: https://docs.pytorch.org/docs/2.13/generated/torch.ao.quantization.quantize_fx.prepare_qat_fx.html

This example demonstrates defining a model, configuring the quantization mapping, and preparing the model for quantization aware training.

```python
import torch
from torch.ao.quantization import get_default_qat_qconfig_mapping
from torch.ao.quantization.quantize_fx import prepare_qat_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().train()
# (optional, but preferred) load the weights from pretrained model
# float_model.load_weights(...)


# define the training loop for quantization aware training
def train_loop(model, train_data):
    model.train()
    for image, target in data_loader:
        ...


# 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=FakeQuantize.with_args(observer=MinMaxObserver.with_args(dtype=torch.qint8)),
#    weight=FakeQuantize.with_args(observer=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_qat_qconfig_mapping("fbgemm")

# We can customize qconfig_mapping in different ways, please take a look at
# the docstring for :func:`~torch.ao.quantization.prepare_fx` for different ways
# to configure this

# 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_qat_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 fake_quantize modules according to the qconfig_mapping
# otherwise the configuration in qconfig_mapping will be ignored
# see :func:`~torch.ao.quantization.prepare_fx` for a detailed explanation of
# how qconfig_mapping interacts with backend_config
prepared_model = prepare_qat_fx(float_model, qconfig_mapping, example_inputs)
# Run training
train_loop(prepared_model, train_loop)
```

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

### General usage examples

Source: https://docs.pytorch.org/docs/2.13/backends.html

Common patterns for querying DSLs and managing operation states.

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

# Query available DSLs
print(pn.available_dsls)  # ['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')
```

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

### Minifier Output Example

Source: https://docs.pytorch.org/docs/2.13/_sources/user_guide/torch_compiler/torch.compiler_aot_inductor_minifier.md.txt

Example output generated when running the minifier with the 'minify' command.

```text
...
W1031 16:48:08.938000 3598491 torch/_dynamo/repro/aoti.py:89] Writing checkpoint with 3 nodes to /data/users/shangdiy/pytorch/torch_compile_debug/run_2024_10_31_16_48_02_720863-pid_3598491/minifier/checkpoints/3.py
W1031 16:48:08.975000 3598491 torch/_dynamo/repro/aoti.py:101] Copying repro file for convenience to /data/users/shangdiy/pytorch/repro.py
Wrote minimal repro out to repro.py
```

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

### Prepare a model for quantization with FX

Source: https://docs.pytorch.org/docs/2.13/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 for calibration.

```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)
```

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

### Initialize FSDP with custom parameter initialization

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

Demonstrates initializing a model on the meta device using a custom initialization function or torchdistX's 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)
```

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

### Example output of Mode logging

Source: https://docs.pytorch.org/docs/2.13/notes/extending.html

The expected console output when running the logging mode examples.

```text
TorchFunctionMode logging:
Function Log: torch.rand(*(10,), **{'requires_grad': True})
Function Log: torch.Tensor.mul(*(tensor([0.7164, 0.9897, 0.1745, 0.9336, 0.4287, 0.7989, 0.2169, 0.7474, 0.5624,
        0.5970], requires_grad=True), 2), **None)
Function Log: torch.Tensor.sum(*(tensor([1.4328, 1.9794, 0.3490, 1.8671, 0.8573, 1.5977, 0.4338, 1.4948, 1.1249,
        1.1939], grad_fn=<MulBackward0>),), **None)
# Note that at the python level, we only see the call to backward but not what happens in the autograd engine.
Function Log: torch.Tensor.backward(*(tensor(12.3307, grad_fn=<SumBackward0>),), **{'gradient': None, 'retain_graph': None, 'create_graph': False, 'inputs': None})

TorchDispatchMode logging:
# Here the requires_grad flag from autograd is removed while default arguments were populated.
Dispatch Log: aten.rand.default(*([10],), **{'device': device(type='cpu'), 'pin_memory': False})
Dispatch Log: aten.mul.Tensor(*(tensor([0.2151, 0.6018, 0.8415, 0.9060, 0.2974, 0.7708, 0.6668, 0.0352, 0.7948,
        0.6023], requires_grad=True), 2), **{})
Dispatch Log: aten.sum.default(*(tensor([0.4303, 1.2036, 1.6831, 1.8120, 0.5949, 1.5416, 1.3335, 0.0705, 1.5897,
        1.2046], grad_fn=<MulBackward0>),), **{})
# Here we don't see the call to backward itself, but its constituents. Starting here with the factory function that creates the initial gradient.
Dispatch Log: aten.ones_like.default(*(tensor(11.4637, grad_fn=<SumBackward0>),), **{'pin_memory': False, 'memory_format': torch.preserve_format})
# This is the backward of the sum
Dispatch Log: aten.expand.default(*(tensor(1.), [10]), **{})
Dispatch Log: aten.mul.Tensor(*(tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]), 2), **{})
Dispatch Log: aten.detach.default(*(tensor([2., 2., 2., 2., 2., 2., 2., 2., 2., 2.]),), **{})
Dispatch Log: aten.detach.default(*(tensor([2., 2., 2., 2., 2., 2., 2., 2., 2., 2.]),), **{})
```

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

### Example Input and Output for get_custom_module_class_keys

Source: https://docs.pytorch.org/docs/2.13/generated/torch.ao.quantization.fx.utils.get_custom_module_class_keys.html

Demonstrates the structure of the input configuration dictionary and the resulting list of unique custom module keys.

```python
{
    QuantType.STATIC: {CustomModule1: ObservedCustomModule},
    QuantType.DYNAMIC: {CustomModule2: DynamicObservedCustomModule},
    QuantType.WEIGHT_ONLY: {CustomModule3: WeightOnlyObservedCustomModule},
}
```

```python
# extract the keys across all inner STATIC, DYNAMIC, and WEIGHT_ONLY dicts
[CustomModule1, CustomModule2, CustomModule3]
```

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

### while_loop usage examples

Source: https://docs.pytorch.org/docs/2.13/higher_order_ops/while_loop.html

Examples demonstrating the use of while_loop with tensor-based and integer-based iteration conditions.

```python
def cond_fn(iter, x):
    return iter.sum() < 10


def body_fn(iter, x):
    return iter + 1, x.sin()


while_loop(cond_fn, body_fn, (torch.zeros(1), torch.randn(3, 4)))
```

```python
def cond_fn(int_iter, x):
    return 2 * int_iter < x.shape[0]


def body_fn(int_iter, x):
    return int_iter + 1, x + int_iter


while_loop(cond_fn, body_fn, (0, torch.randn(3, 4)))
```

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

### Initialize and step SGD optimizer

Source: https://docs.pytorch.org/docs/2.13/generated/torch.optim.SGD.html

Demonstrates the standard workflow of initializing the SGD optimizer, clearing gradients, computing the loss, and performing an optimization step.

```python
>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> optimizer.zero_grad()
>>> loss_fn(model(input), target).backward()
>>> optimizer.step()
```

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

### GaussianNLLLoss Usage Examples

Source: https://docs.pytorch.org/docs/2.13/generated/torch.nn.modules.loss.GaussianNLLLoss.html

Examples demonstrating the use of GaussianNLLLoss with heteroscedastic and homoscedastic variance inputs.

```python
>>> loss = nn.GaussianNLLLoss()
>>> input = torch.randn(5, 2, requires_grad=True)
>>> target = torch.randn(5, 2)
>>> var = torch.ones(5, 2, requires_grad=True)  # heteroscedastic
>>> output = loss(input, target, var)
>>> output.backward()
```

```python
>>> loss = nn.GaussianNLLLoss()
>>> input = torch.randn(5, 2, requires_grad=True)
>>> target = torch.randn(5, 2)
>>> var = torch.ones(5, 1, requires_grad=True)  # homoscedastic
>>> output = loss(input, target, var)
>>> output.backward()
```

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

### Implement example_arguments

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

Define sample inputs for tracing and testing by returning positional and keyword arguments.

```python
def example_arguments(self):
    # For a model expecting a single tensor input
    return (torch.randn(1, 3, 224, 224),), None


def example_arguments(self):
    # For a model with multiple inputs and keyword arguments
    return (torch.randn(1, 3, 224, 224), torch.randn(1, 512)), {
        "temperature": 1.0
    }
```

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

### DTensor Sharding Configuration Example

Source: https://docs.pytorch.org/docs/2.13/distributed.tensor.html

Example configuration for a 2D mesh with default sharding behavior.

```text
tensor shape: [8, 8]
mesh: [[0, 1], [2, 3]], names=("dp", "tp")
placements: [Shard(0), Shard(0)]
```

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

### Data Dependent Operation Example

Source: https://docs.pytorch.org/docs/2.13/_sources/nested.md.txt

Example of code that triggers a data-dependent operation error during compilation.

```python
>>> a = torch.randn(50, 128)
>>> b = torch.randn(32, 128)
>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged, dtype=torch.float32)
>>> def f(nt): return nt.chunk(2, dim=0)[0]
...
>>> compiled_f = torch.compile(f, fullgraph=True)
>>> output = compiled_f(nt)
```

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

### example_arguments()

Source: https://docs.pytorch.org/docs/2.13/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 representing the expected input shapes and types for the model's forward method.

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

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

### Initializing and using Conv3d

Source: https://docs.pytorch.org/docs/2.13/generated/torch.ao.nn.quantized.dynamic.modules.conv.Conv3d.html

Examples showing how to initialize the Conv3d module with various kernel, stride, padding, and dilation configurations, followed by a forward pass with a random input tensor.

```python
>>> # With square kernels and equal stride
>>> m = nn.quantized.dynamic.Conv3d(16, 33, 3, stride=2)
>>> # non-square kernels and unequal stride and with padding
>>> m = nn.quantized.dynamic.Conv3d(16, 33, (3, 5, 5), stride=(1, 2, 2), padding=(1, 2, 2))
>>> # non-square kernels and unequal stride and with padding and dilation
>>> m = nn.quantized.dynamic.Conv3d(16, 33, (3, 5, 5), stride=(1, 2, 2), padding=(1, 2, 2), dilation=(1, 2, 2))
>>> input = torch.randn(20, 16, 56, 56, 56)
>>> output = m(input)
```

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

### Materialized LOWER_RIGHT Bias Example

Source: https://docs.pytorch.org/docs/2.13/generated/torch.nn.attention.bias.CausalVariant.html

Example output of a LOWER_RIGHT causal bias tensor with shape (3,4).

```text
[[1, 1, 0, 0],
 [1, 1, 1, 0],
 [1, 1, 1, 1]]
```

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

### Create a C++ Extension with setuptools

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

Uses CppExtension to define a C++ extension module within a setup.py script.

```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
...     })
```

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

### Materialized UPPER_LEFT Bias Example

Source: https://docs.pytorch.org/docs/2.13/generated/torch.nn.attention.bias.CausalVariant.html

Example output of an UPPER_LEFT causal bias tensor with shape (3,4).

```text
[[1, 0, 0, 0],
 [1, 1, 0, 0],
 [1, 1, 1, 0]]
```

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

### Pickle GLOBAL opcode example

Source: https://docs.pytorch.org/docs/2.13/_sources/package.md.txt

Example of a GLOBAL opcode used by pickle to identify object types.

```text
GLOBAL 'torchvision.models.resnet Resnet`
```

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

### Sample from LKJCholesky

Source: https://docs.pytorch.org/docs/2.13/distributions.html

Demonstrates initializing the distribution and generating a sample.

```python
>>> l = LKJCholesky(3, 0.5)
>>> l.sample()  # l @ l.T is a sample of a correlation 3x3 matrix
tensor([[ 1.0000,  0.0000,  0.0000],
        [ 0.3516,  0.9361,  0.0000],
        [-0.1899,  0.4748,  0.8593]])
```

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

### Initialize SummaryWriter

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

Demonstrates different 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/
```

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

### Install ONNX Exporter Dependencies

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

Install the required ONNX and ONNX Script packages via pip.

```bash
pip install --upgrade onnx onnxscript
```

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

### Define a SYCL extension with setup

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

Use SyclExtension to define an extension module and BuildExtension to handle the build process in setup.py.

```python
>>> from torch.utils.cpp_extension import BuildExtension, SyclExtension
>>> setup(
...     name='xpu_extension',
...     ext_modules=[
...     SyclExtension(
...                 name='xpu_extension',
...                 sources=['extension.cpp', 'extension_kernel.cpp'],
...                 extra_compile_args={'cxx': ['-g', '-std=c++20', '-fPIC']})
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })
```

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

### torch.arange(start=0, end, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False)

Source: https://docs.pytorch.org/docs/2.13/generated/torch.arange.html

Returns a 1-D tensor of size ceil((end - start) / step) with values from the interval [start, end) taken with common difference step beginning from start.

```APIDOC
## torch.arange

### Description
Returns a 1-D tensor of size ceil((end - start) / step) with values from the interval [start, end) taken with common difference step beginning from start.

### Parameters
- **start** (Number, optional) - The starting value for the set of points. Default: 0.
- **end** (Number) - The ending value for the set of points.
- **step** (Number, optional) - The gap between each pair of adjacent points. Default: 1.

### Keyword Arguments
- **out** (Tensor, optional) - The output tensor.
- **dtype** (torch.dtype, optional) - The desired data type of returned tensor.
- **layout** (torch.layout, optional) - The desired layout of returned Tensor. Default: torch.strided.
- **device** (torch.device, optional) - The desired device of returned tensor.
- **requires_grad** (bool, optional) - If autograd should record operations on the returned tensor. Default: False.

### Example
```python
>>> torch.arange(5)
tensor([ 0,  1,  2,  3,  4])
>>> torch.arange(1, 4)
tensor([ 1,  2,  3])
>>> torch.arange(1, 2.5, 0.5)
tensor([ 1.0000,  1.5000,  2.0000])
```
```

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

### Initialize and use SyncBatchNorm

Source: https://docs.pytorch.org/docs/2.13/generated/torch.nn.SyncBatchNorm.html

Demonstrates initializing SyncBatchNorm with learnable parameters and using a custom process group for distributed operations.

```python
>>> # With Learnable Parameters
>>> m = nn.SyncBatchNorm(100)
>>> # creating process group (optional)
>>> # ranks is a list of int identifying rank ids.
>>> ranks = list(range(8))
>>> r1, r2 = ranks[:4], ranks[4:]
>>> # Note: every rank calls into new_group for every
>>> # process group created, even if that rank is not
>>> # part of the group.
>>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]
>>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]
>>> # Without Learnable Parameters
>>> m = nn.BatchNorm3d(100, affine=False, process_group=process_group)
>>> input = torch.randn(20, 100, 35, 45, 10)
>>> output = m(input)

>>> # network is nn.BatchNorm layer
>>> sync_bn_network = nn.SyncBatchNorm.convert_sync_batchnorm(network, process_group)
>>> # only single gpu per process is currently supported
>>> ddp_sync_bn_network = torch.nn.parallel.DistributedDataParallel(
>>>                         sync_bn_network,
>>>                         device_ids=[args.local_rank],
>>>                         output_device=args.local_rank)
```

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

### Trace Log Example

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

Example of a trace log output generated when running with TORCH_LOGS=dynamo, showing the stack state.

```text
TRACE LOAD_GLOBAL y [TorchInGraphFunctionVariable(<built-in method any>), TensorVariable()]
```

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

### Implement a custom autograd function with setup_context

Source: https://docs.pytorch.org/docs/2.13/_sources/notes/extending.md.txt

Shows how to save intermediate values for backward passes by returning them as outputs and using setup_context to manage the context.

```python
    class MyCube(torch.autograd.Function):
        @staticmethod
        def forward(x):
            # We wish to save dx for backward. In order to do so, it must
            # be returned as an output.
            dx = 3 * x ** 2
            result = x ** 3
            return result, dx

        @staticmethod
        def setup_context(ctx, inputs, output):
            x, = inputs
            result, dx = output
            ctx.save_for_backward(x, dx)

        @staticmethod
        def backward(ctx, grad_output, grad_dx):
            x, dx = ctx.saved_tensors
            # In order for the autograd.Function to work with higher-order
            # gradients, we must add the gradient contribution of `dx`,
            # which is grad_dx * 6 * x.
            result = grad_output * dx + grad_dx * 6 * x
            return result

    # Wrap MyCube in a function so that it is clearer what the output is
    def my_cube(x):
        result, dx = MyCube.apply(x)
        return result
```

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

### find_closure_group(input_string, start, group)

Source: https://docs.pytorch.org/docs/2.13/generated/torch.utils.hipify.hipify_python.find_closure_group.html

Finds the first balanced closure group in a string starting from a specific position.

```APIDOC
## find_closure_group(input_string, start, group)

### Description
Finds the positions of the first balanced closure group (e.g., parentheses or brackets) in an input string starting from a given index.

### Parameters
- **input_string** (str) - The string to search within.
- **start** (int) - The starting index in the input string.
- **group** (list) - A list containing the opening and closing characters, e.g., ['(', ')'].

### Returns
- **tuple** - A tuple containing the indices of the opening and closing characters of the found group.

### Example
```python
>>> find_closure_group("(hi)", 0, ["(", ")"])
(0, 3)
```
```

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

### Configure QConfig with custom observer arguments

Source: https://docs.pytorch.org/docs/2.13/generated/torch.ao.quantization.qconfig.QConfig.html

Demonstrates how to instantiate a QConfig object using observer classes configured with specific data types via the with_args method.

```python
my_qconfig = QConfig(
    activation=MinMaxObserver.with_args(dtype=torch.qint8),
    weight=default_observer.with_args(dtype=torch.qint8),
)
```

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

### extract_arguments(start, string)

Source: https://docs.pytorch.org/docs/2.13/generated/torch.utils.hipify.hipify_python.extract_arguments.html

Parses a function parameter closure string to extract the start and end indices of each argument.

```APIDOC
## extract_arguments(start, string)

### Description
Returns the list of arguments in the upcoming function parameter closure.

### Parameters
- **start** (int) - The starting index for the extraction process.
- **string** (str) - The input string containing the function parameter closure (e.g., '(blocks, threads, 0, THCState_getCurrentStream(state))').

### Returns
- **list** - A list of dictionaries, where each dictionary contains 'start' and 'end' keys representing the character indices of an argument.
```

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

### Initialize and Use HashStore

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

Shows how to initialize a thread-safe HashStore for use within a single process.

```python
>>> import torch.distributed as dist
>>> store = dist.HashStore()
>>> # store can be used from other threads
>>> # Use any of the store methods after initialization
>>> store.set("first_key", "first_value")
```

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

### Numpy dtypes Unpickling Error Examples

Source: https://docs.pytorch.org/docs/2.13/notes/serialization.html

Examples of errors encountered when numpy dtypes are not properly allowlisted during weights_only loading.

```text
WeightsUnpickler error: Can only build Tensor, Parameter, OrderedDict or types allowlisted via `add_safe_globals`,
but got <class 'numpy.dtype[float32]'>
```

```text
WeightsUnpickler error: Can only build Tensor, Parameter, OrderedDict or types allowlisted via `add_safe_globals`,
but got <class 'numpy.dtypes.Float32DType'>
```

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

### torch.Tensor.narrow

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

Returns a new tensor that is a narrowed version of the input tensor. The dimension dim is narrowed from start to start + length.

```APIDOC
## torch.Tensor.narrow

### Description
Returns a new tensor that is a narrowed version of the input tensor. The dimension `dim` is narrowed from `start` to `start + length`.

### Signature
`torch.Tensor.narrow(dim, start, length)`

### Parameters
- **dim** (int) - The dimension along which to narrow.
- **start** (int) - The starting index.
- **length** (int) - The distance to the ending index.
```

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

### Configure PrepareCustomConfig

Source: https://docs.pytorch.org/docs/2.13/generated/torch.ao.quantization.fx.custom_config.PrepareCustomConfig.html

Demonstrates the fluent API usage for setting various quantization configurations such as standalone modules, non-traceable modules, and attribute preservation.

```python
prepare_custom_config = PrepareCustomConfig()             .set_standalone_module_name("module1", qconfig_mapping, example_inputs,                 child_prepare_custom_config, backend_config)             .set_standalone_module_class(MyStandaloneModule, qconfig_mapping, example_inputs,                 child_prepare_custom_config, backend_config)             .set_float_to_observed_mapping(FloatCustomModule, ObservedCustomModule)             .set_non_traceable_module_names(["module2", "module3"])             .set_non_traceable_module_classes([NonTraceableModule1, NonTraceableModule2])             .set_input_quantized_indexes([0])             .set_output_quantized_indexes([0])             .set_preserved_attributes(["attr1", "attr2"])
```

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

### Configure FX Conversion with ConvertCustomConfig

Source: https://docs.pytorch.org/docs/2.13/generated/torch.ao.quantization.fx.custom_config.ConvertCustomConfig.html

Demonstrates initializing a configuration object and chaining methods to set module mappings and preserved attributes.

```python
convert_custom_config = ConvertCustomConfig()             .set_observed_to_quantized_mapping(ObservedCustomModule, QuantizedCustomModule)             .set_preserved_attributes(["attr1", "attr2"])
```

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

### Standardized Callgrind Output Example

Source: https://docs.pytorch.org/docs/2.13/benchmark_utils.html

Example of raw Callgrind output showing path prefixes that can cause issues when diffing profiles.

```text
23234231 /tmp/first_build_dir/thing.c:foo(...)
 9823794 /tmp/first_build_dir/thing.c:bar(...)
  ...
   53453 .../aten/src/Aten/...:function_that_actually_changed(...)
  ...
 -9823794 /tmp/second_build_dir/thing.c:bar(...)
-23234231 /tmp/second_build_dir/thing.c:foo(...)
```

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

### Fractional Max Pooling Examples

Source: https://docs.pytorch.org/docs/2.13/generated/torch.nn.functional.fractional_max_pool2d.html

Demonstrates applying fractional max pooling using both a fixed output size and an output ratio.

```python
>>> input = torch.randn(20, 16, 50, 32)
>>> # pool of square window of size=3, and target output size 13x12
>>> F.fractional_max_pool2d(input, 3, output_size=(13, 12))
>>> # pool of square window and target output size being half of input image size
>>> F.fractional_max_pool2d(input, 3, output_ratio=(0.5, 0.5))
```

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

### Install PyTorch for Intel GPU (XPU)

Source: https://docs.pytorch.org/docs/2.13/_sources/notes/get_start_xpu.md.txt

Commands to install PyTorch, torchvision, and torchaudio for Intel GPU support via pip.

```bash
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/xpu
```

```bash
pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/xpu
```

```bash
pip install torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/xpu
```

```bash
pip install torch==2.10.0 torchvision==0.25.0 torchaudio==2.10.0 --index-url https://download.pytorch.org/whl/xpu
```

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

### torch.cuda.nvtx.range_start(msg)

Source: https://docs.pytorch.org/docs/2.13/_sources/generated/torch.cuda.nvtx.range_start.rst.txt

Starts a nested range in the NVTX profiling tool. This function is used to define the start of a labeled range for performance analysis.

```APIDOC
## torch.cuda.nvtx.range_start

### Description
Starts an NVTX range with the specified message. This is used to instrument code for performance profiling in NVIDIA tools.

### Signature
torch.cuda.nvtx.range_start(msg: str) -> int

### Parameters
- **msg** (str) - The message or label to associate with the start of the range.

### Returns
- **int** - A handle representing the range, which should be passed to `torch.cuda.nvtx.range_end` to close the range.
```

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

### Implement a model entrypoint

Source: https://docs.pytorch.org/docs/2.13/_sources/hub.md.txt

Example implementation of a resnet18 entrypoint, including dependency declaration and docstring for help messages.

```python
  dependencies = ['torch']
  from torchvision.models.resnet import resnet18 as _resnet18

  # resnet18 is the name of entrypoint
  def resnet18(pretrained=False, **kwargs):
      """ # This docstring shows up in hub.help()
      Resnet18 model
      pretrained (bool): kwargs, load pretrained weights into the model
      """
      # Call the model, load pretrained weights
      model = _resnet18(pretrained=pretrained, **kwargs)
      return model
```

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

### FractionalMaxPool3d Usage Examples

Source: https://docs.pytorch.org/docs/2.13/generated/torch.nn.FractionalMaxPool3d.html

Demonstrates initializing FractionalMaxPool3d with either a fixed output size or an output ratio, followed by a forward pass with a random input tensor.

```python
>>> # pool of cubic window of size=3, and target output size 13x12x11
>>> m = nn.FractionalMaxPool3d(3, output_size=(13, 12, 11))
>>> # pool of cubic window and target output size being half of input size
>>> m = nn.FractionalMaxPool3d(3, output_ratio=(0.5, 0.5, 0.5))
>>> input = torch.randn(20, 16, 50, 32, 16)
>>> output = m(input)
```
Tokens:
Raw