PyTorch

https://docs.pytorch.org/docs/2.14
Admin
PyTorch is an optimized tensor library for deep learning using GPUs and CPUs.
Tokens:3,523,146
Snippets:26,885
Trust Score:10
Update:1 week ago
Show doc for...
Tokens:
Raw
### Quickstart

Source: https://docs.pytorch.org/docs/2.14/nn.attention.flex_attention.html

Quickstart guide for elastic training.

```APIDOC
## Quickstart

### Description
Quickstart guide for elastic training.

### Method
N/A (module)

### Endpoint
N/A

### Parameters
None

### Request Example
N/A

### Response
N/A
```

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

### Quickstart

Source: https://docs.pytorch.org/docs/2.14/future_mod.html

Quickstart guide for torch.distributed.elastic.

```APIDOC
## Quickstart

### Description
Quickstart guide for torch.distributed.elastic.

### Method
N/A (Concept)

### Endpoint
N/A (Python module)

### Parameters
None

### Request Example
N/A

### Response
N/A

### Error Handling
N/A
```

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

### SyclExtension setup example

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

Example of using SyclExtension and BuildExtension in a setup.py to build a SYCL extension. Requires Ninja and specifies compile arguments for C++.

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

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

### EtcdServer usage example

Source: https://docs.pytorch.org/docs/2.14/elastic/rendezvous.html

Shows how to instantiate EtcdServer with a binary path, port, and data directory, start it, obtain a client, and stop the server. The termination handler is not a substitute for calling stop().

```python
server = EtcdServer("/usr/bin/etcd", 2379, "/tmp/default.etcd")
server.start()
client = server.get_client()
# use client
server.stop()
```

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

### torch.quantization.get_fqn_to_example_inputs

Source: https://docs.pytorch.org/docs/2.14/utils.html

Gets a mapping from fully qualified names to example inputs. This is part of the quantization API.

```APIDOC
## torch.quantization.get_fqn_to_example_inputs

### Description
Gets a mapping from fully qualified names to example inputs. This is part of the quantization API.

### Method
get_fqn_to_example_inputs

### Parameters
- **model** (Module) - Required - The model.
- **example_inputs** (tuple) - Required - Example inputs.

### Returns
- **fqn_to_example_inputs** (dict) - Mapping from FQN to example inputs.
```

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

### Example: prepare_qat_fx usage

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

This example demonstrates how to use prepare_qat_fx to prepare a model for quantization aware training. It defines a model with submodules, creates a qconfig_mapping using get_default_qat_qconfig_mapping, and calls prepare_qat_fx with example inputs. The example also includes a placeholder training loop and comments on customizing qconfig.

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

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

### TCPStore server and client example

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

Shows how to create a TCPStore server and client, then use store methods. The server is initialized with is_master=True and a timeout of 30 seconds; the client uses the same host, port, and world_size. After initialization, both can call store methods like set() and get().

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

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

### get_fqn_to_example_inputs

Source: https://docs.pytorch.org/docs/2.14/quantization.html

Gets a mapping from fully qualified names to example inputs. Used for FX graph quantization.

```APIDOC
## get_fqn_to_example_inputs

### Description
Gets a mapping from fully qualified names to example inputs. Used for FX graph quantization.

### Method
N/A (Function call)

### Endpoint
N/A (Function call)

### Parameters
- **model** (torch.nn.Module) - Required - The model.
- **example_inputs** (tuple) - Required - Example inputs.

### Request Example
N/A (Function call)

### Response
#### Success Response
- **fqn_to_example_inputs** (dict) - Mapping from FQN to example inputs.

#### Response Example
N/A (Function call)
```

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

### FileStore example

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

Creates two FileStore instances pointing to the same file, demonstrating that multiple processes can use the same file-based store. After initialization, store methods like set() and get() can be used from either instance.

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

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

### Examples

Source: https://docs.pytorch.org/docs/2.14/nn.attention.flex_attention.html

Examples for elastic training.

```APIDOC
## Examples

### Description
Examples for elastic training.

### Method
N/A (module)

### Endpoint
N/A

### Parameters
None

### Request Example
N/A

### Response
N/A
```

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

### Examples

Source: https://docs.pytorch.org/docs/2.14/future_mod.html

Examples for torch.distributed.elastic.

```APIDOC
## Examples

### Description
Examples for torch.distributed.elastic.

### Method
N/A (Concept)

### Endpoint
N/A (Python module)

### Parameters
None

### Request Example
N/A

### Response
N/A

### Error Handling
N/A
```

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

### prepare_fx example: model definition, calibration, and preparation

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

This example demonstrates how to use prepare_fx to insert observers into a model for quantization. It defines a model with a submodule, a calibration function, and shows how to create a qconfig_mapping using get_default_qconfig_mapping. The example also includes commented-out customizations for qconfig and qconfig_mapping, and ends with the prepare_fx call and 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 user 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)

```

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

### torch.logspace examples

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

Examples of using torch.logspace with different start, end, steps, and base values. The first example shows the default base of 10.0, the second uses non-integer start and end, the third uses steps=1, and the fourth uses base=2.

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

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

### torch.lerp example with scalar and tensor weight

Source: https://docs.pytorch.org/docs/2.14/generated/torch.lerp.html

Demonstrates torch.lerp with a scalar weight (0.5) and a tensor weight (torch.full_like(start, 0.5)). The example creates start and end tensors, prints them, and shows the interpolated output for both weight types.

```python
>>> start = torch.arange(1., 5.)
>>> end = torch.empty(4).fill_(10)
>>> start
tensor([ 1.,  2.,  3.,  4.])
>>> end
tensor([ 10.,  10.,  10.,  10.])
>>> torch.lerp(start, end, 0.5)
tensor([ 5.5000,  6.0000,  6.5000,  7.0000])
>>> torch.lerp(start, end, torch.full_like(start, 0.5))
tensor([ 5.5000,  6.0000,  6.5000,  7.0000])
```

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

### Example: Using Tensor.get_device() with CUDA and CPU tensors

Source: https://docs.pytorch.org/docs/2.14/generated/torch.Tensor.get_device.html

Shows how to call get_device() on a CUDA tensor to get its device ordinal (0) and on a CPU tensor to get -1. The example creates a tensor on 'cuda:0' and then moves it to CPU with .cpu().

```python
>>> x = torch.randn(3, 4, 5, device='cuda:0')
>>> x.get_device()
0
>>> x.cpu().get_device()
-1
```

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

### Create model, optimizer, and state dict, then create FileSystemReader

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

Shows the initial setup for loading a checkpoint: creating a model, optimizer, and state dict, and constructing a FileSystemReader for the checkpoint path.

```python
>>> my_model = MyModule()
>>> optimizer = Adagrad(my_model.parameters())
>>> model_state_dict = my_model.state_dict()
>>> fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(
...     "/checkpoint/1"
... )

```

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

### Example: torch.Tensor.size()

Source: https://docs.pytorch.org/docs/2.14/generated/torch.Tensor.size.html

Shows how to call size() without dim to get a torch.Size, and with dim=1 to get the size of a specific dimension.

```python
>>> t = torch.empty(3, 4, 5)
>>> t.size()
torch.Size([3, 4, 5])
>>> t.size(dim=1)
4
```

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

### _static _setup_context(_ctx_, _inputs_, _output_)

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

Sets up the context for backward when forward is defined with the signature forward(*args, **kwargs). Override this method to set up ctx for backward.

```APIDOC
## _static _setup_context(_ctx_, _inputs_, _output_)

### Description
There are two ways to define the forward pass of an autograd.Function. Either:
1. Override forward with the signature `forward(ctx, *args, **kwargs)`. `setup_context` is not overridden. Setting up the ctx for backward happens inside the `forward`.
2. Override forward with the signature `forward(*args, **kwargs)` and override `setup_context`. Setting up the ctx for backward happens inside `setup_context` (as opposed to inside the `forward`).

See `torch.autograd.Function.forward()` and Extending torch.autograd for more details.

### Parameters
- **ctx** (torch.autograd.function._ContextMethodMixin) - Required - The context object to set up.
- **inputs** (Any) - Required - The inputs to the forward function.
- **output** (Any) - Required - The output of the forward function.

### Return Type
- **Any**
```

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

### torch.narrow example

Source: https://docs.pytorch.org/docs/2.14/generated/torch.narrow.html

Demonstrates narrowing a 2D tensor along dimension 0 and dimension 1, and using a negative start index with a tensor start value.

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

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

### StorageWriter._set_up_storage_writer

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

Initialize this instance.

```APIDOC
## StorageWriter._set_up_storage_writer

### Description
Initialize this instance.

### Method
_set_up_storage_writer

### Signature
`_set_up_storage_writer(is_coordinator, *args, **kwargs)`

### Parameters
- **is_coordinator** (bool) - Required - Whether this instance is responsible for coordinating the checkpoint.
```

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

### start_processes examples: function and binary entrypoints

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

Shows valid and invalid calls to start_processes. The first call launches two copies of function foo with args ("bar0") and ("bar1"). The second is invalid because envs is missing for local rank 1. The third launches two copies of /usr/bin/touch. The fourth launches /usr/bin/echo, where arguments are cast to strings, so it runs echo "1" "2" "3" and echo "[1, 2, 3]".

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

```

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

### _setup_context(ctx, inputs, output)

Source: https://docs.pytorch.org/docs/2.14/generated/torch.autograd.function.InplaceFunction.html

Sets up the context for backward when using the second forward signature style.

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

### Description
There are two ways to define the forward pass of an autograd.Function. Either override forward with the signature `forward(ctx, *args, **kwargs)` and do not override `setup_context`, or override forward with the signature `forward(*args, **kwargs)` and override `setup_context`. Setting up the ctx for backward happens inside `setup_context` (as opposed to inside the `forward`).

### Method
_setup_context

### Parameters
- **ctx** (Any) - Required - The context object to set up for backward.
- **inputs** (Any) - Required - The inputs to the forward pass.
- **output** (Any) - Required - The output of the forward pass.

### Returns
- **Any** - The return type is `Any`.
```

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

### Example: Using torch.func.linearize

Source: https://docs.pytorch.org/docs/2.14/generated/torch.func.linearize.html

Shows how to use torch.func.linearize with a simple function. The example defines a function that applies sin, calls linearize with a zero tensor, and then uses the returned jvp_fn with a ones tensor to get the derivative (cos(0)=1).

```python
>>> import torch
>>> from torch.func import linearize
>>> def fn(x):
...     return x.sin()
...
>>> output, jvp_fn = linearize(fn, torch.zeros(3, 3))
>>> jvp_fn(torch.ones(3, 3))
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]])
>>>
```

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

### CppExtension

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

Creates a setuptools.Extension for C++ with minimal arguments. All arguments are forwarded to the setuptools.Extension constructor.

```APIDOC
## CppExtension

### Description
Creates a `setuptools.Extension` for C++ with the bare minimum arguments to build a C++ extension. All arguments are forwarded to the `setuptools.Extension` constructor.

### Method
Constructor

### Signature
`torch.utils.cpp_extension.CppExtension(name, sources, *args, **kwargs)`

### Parameters
- **name** (str) - Required - The name of the extension.
- **sources** (list[str]) - Required - List of source file paths.
- **\*args** - Optional - Additional positional arguments forwarded to `setuptools.Extension`.
- **\*\*kwargs** - Optional - Additional keyword arguments forwarded to `setuptools.Extension`.

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

### Notes
- The PyTorch python API cannot be built with `py_limited_api=True`. If used, the user must avoid libtorch_python APIs and only use libtorch APIs.
- PyTorch defines `-DPy_LIMITED_API=min_supported_cpython` when `py_limited_api` is specified.
```

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

### Example: Using torch.Tensor.item()

Source: https://docs.pytorch.org/docs/2.14/generated/torch.Tensor.item.html

Shows how to call item() on a single-element tensor to get a Python number. The example creates a tensor with one element (1.0) and prints the result 1.0. Note that item() only works for tensors with one element and is not differentiable.

```python
>>> x = torch.tensor([1.0])
>>> x.item()
1.0
```

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

### PrepareCustomConfig example usage

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

Shows how to build a PrepareCustomConfig with chained setter methods for standalone modules, float-to-observed mapping, non-traceable modules, quantized indexes, and preserved attributes.

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

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

### Example: itertools.product vs torch.cartesian_prod

Source: https://docs.pytorch.org/docs/2.14/generated/torch.cartesian_prod.html

Shows the equivalent behavior of Python's itertools.product and torch.cartesian_prod on the same input lists. The example first creates lists a and b, then demonstrates itertools.product, and finally converts them to tensors and calls torch.cartesian_prod to get the same result as a tensor.

```python
>>> import itertools
>>> a = [1, 2, 3]
>>> b = [4, 5]
>>> list(itertools.product(a, b))
[(1, 4), (1, 5), (2, 4), (2, 5), (3, 4), (3, 5)]
>>> tensor_a = torch.tensor(a)
>>> tensor_b = torch.tensor(b)
>>> torch.cartesian_prod(tensor_a, tensor_b)
tensor([[1, 4],
        [1, 5],
        [2, 4],
        [2, 5],
        [3, 4],
        [3, 5]])
```

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

### CppExtension setup.py example

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

Example of using CppExtension to create a C++ extension. It imports BuildExtension and CppExtension, and passes extra_compile_args and extra_link_args. Note the warning about py_limited_api.

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

```

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

### torch.linspace examples

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

Shows basic usage of torch.linspace with different start, end, and steps values, including keyword arguments and a single-step case.

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

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

### Use openreg device with explicit import

Source: https://docs.pytorch.org/docs/2.14/accelerator/autoload.html

After building and installing the backend, this example shows using the openreg device with an explicit import of torch_openreg.

```python
>>> import torch
>>> import torch_openreg
>>> torch.tensor(1, device="openreg")
tensor(1, device='openreg:0')
```

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

### Example: distribute_module with partition_fn, input_fn, and output_fn

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

Shows how to use distribute_module to convert a module's parameters to DTensors, replicate inputs, and convert outputs back to plain tensors. The example defines a MyModule, initializes a device mesh, and provides callbacks for partitioning, input conversion, and output conversion.

```python
import torch.nn as nn
from torch.distributed.tensor import (
    DTensor,
    Replicate,
    Shard,
    distribute_module,
    distribute_tensor,
    init_device_mesh,
)


class MyModule(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.fc = nn.Linear(8, 8)

    def forward(self, x):
        return self.fc(x)


mesh = init_device_mesh("cuda", (4,))


def shard_params(name, mod, mesh):
    if isinstance(mod, nn.Linear):
        for pname, param in mod.named_parameters(recurse=False):
            mod.register_parameter(
                pname, nn.Parameter(distribute_tensor(param, mesh, [Shard(0)]))
            )


# Replicate the input across the mesh before forward.
def to_replicated(mod, inputs, mesh):
    return tuple(
        distribute_tensor(x, mesh, [Replicate()])
        if isinstance(x, torch.Tensor)
        else x
        for x in inputs
    )


# Convert the DTensor output back to a plain torch.Tensor after forward.
def to_local(mod, outputs, mesh):
    return outputs.full_tensor() if isinstance(outputs, DTensor) else outputs


sharded = distribute_module(
    MyModule(),
    mesh,
    partition_fn=shard_params,
    input_fn=to_replicated,
    output_fn=to_local,
)

```

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

### Example: Using unpack_dual with dual_level

Source: https://docs.pytorch.org/docs/2.14/generated/torch.autograd.forward_ad.unpack_dual.html

Shows how to use torch.autograd.forward_ad.unpack_dual inside a dual_level context to extract the primal and tangent (JVP) from a dual tensor. The example creates a dual tensor with make_dual, passes it through a function f, and then unpacks the output to get the value and the forward-mode gradient.

```python
>>> with dual_level():
...     inp = make_dual(x, x_t)
...     out = f(inp)
...     y, jvp = unpack_dual(out)
...     jvp = unpack_dual(out).tangent
```

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

### torch.xpu.init

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

Initializes the XPU backend.

```APIDOC
## torch.xpu.init

### Description
Initializes the XPU backend.

### Method
Function call

### Endpoint
`torch.xpu.init()`

### Returns
- None
```

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

### Example usage of torch.accelerator.synchronize

Source: https://docs.pytorch.org/docs/2.14/generated/torch.accelerator.synchronize.html

Shows how to measure elapsed time for a tensor sum operation on the current accelerator using torch.accelerator.synchronize. Requires an available accelerator; otherwise the assertion fails. The example records start and end events, performs a sum, synchronizes, and then computes elapsed time.

```python
>>> assert torch.accelerator.is_available() "No available accelerators detected."
>>> start_event = torch.Event(enable_timing=True)
>>> end_event = torch.Event(enable_timing=True)
>>> start_event.record()
>>> tensor = torch.randn(100, device=torch.accelerator.current_accelerator())
>>> sum = torch.sum(tensor)
>>> end_event.record()
>>> torch.accelerator.synchronize()
>>> elapsed_time_ms = start_event.elapsed_time(end_event)
```

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

### CUDAExtension setup.py example

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

Example of using CUDAExtension to create a CUDA/C++ extension. It imports BuildExtension and CUDAExtension, and passes sources including a .cu file, with separate compile args for cxx and nvcc.

```python
>>> from setuptools import setup
>>> from torch.utils.cpp_extension import BuildExtension, CUDAExtension
>>> setup(
...     name='cuda_extension',
...     ext_modules=[
...         CUDAExtension(
...                 name='cuda_extension',
...                 sources=['extension.cpp', 'extension_kernel.cu'],
...                 extra_compile_args={'cxx': ['-g'],
...                                     'nvcc': ['-O2']},
...                 extra_link_args=['-Wl,--no-as-needed', '-lcuda'])
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })

```

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

### Generate Hermitian-symmetric time-domain signal with torch.fft.ihfft2

Source: https://docs.pytorch.org/docs/2.14/generated/torch.fft.hfft2.html

Starting from a real frequency-space signal T, generate a Hermitian-symmetric time-domain signal t using torch.fft.ihfft2. This is the first step of the example.

```python
>>> T = torch.rand(10, 9)
>>> t = torch.fft.ihfft2(T)
```

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

### Configure a model per stage and create a PipelineStage

Source: https://docs.pytorch.org/docs/2.14/distributed.pipelining.html

Shows how to initialize the whole model on meta-device, delete undesired layers for a given stage, and wrap the resulting model in a PipelineStage. This is a simple 2-stage example.

```python
with torch.device("meta"):
    assert num_stages == 2, "This is a simple 2-stage example"

    # we construct the entire model, then delete the parts we do not need for this stage
    # in practice, this can be done using a helper function that automatically divides up layers across stages.
    model = Transformer()

    if stage_index == 0:
        # prepare the first stage model
        del model.layers["1"]
        model.norm = None
        model.output = None

    elif stage_index == 1:
        # prepare the second stage model
        model.tok_embeddings = None
        del model.layers["0"]

    from torch.distributed.pipelining import PipelineStage
    stage = PipelineStage(
        model,
        stage_index,
        num_stages,
        device,
    )

```

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

### Example: Using torch.neg

Source: https://docs.pytorch.org/docs/2.14/generated/torch.neg.html

Creates a random tensor with torch.randn(5), prints it, then applies torch.neg to get the element-wise negation. The output shows the original and negated tensors.

```python
>>> a = torch.randn(5)
>>> a
tensor([ 0.0090, -0.2262, -0.0682, -0.2866,  0.3940])
>>> torch.neg(a)
tensor([-0.0090,  0.2262,  0.0682,  0.2866, -0.3940])
```

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

### Generate Hermitian-symmetric time-domain signal with torch.fft.ihfftn

Source: https://docs.pytorch.org/docs/2.14/generated/torch.fft.hfftn.html

Creates a real frequency-space tensor and transforms it to a Hermitian-symmetric time-domain signal using torch.fft.ihfftn. This is the starting point for the hfftn round-trip examples.

```python
>>> T = torch.rand(10, 9)
>>> t = torch.fft.ihfftn(T)
```

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

### Example: Create and use a dynamic quantized Linear module

Source: https://docs.pytorch.org/docs/2.14/generated/torch.ao.nn.quantized.dynamic.modules.linear.Linear.html

Shows how to instantiate a dynamic quantized Linear module with 20 input features and 30 output features, pass a random input tensor, and print the output size.

```python
>>> m = nn.quantized.dynamic.Linear(20, 30)
>>> input = torch.randn(128, 20)
>>> output = m(input)
>>> print(output.size())
torch.Size([128, 30])
```

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

### FSDP with param_init_fn and torchdistX deferred_init

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

Shows how to initialize a meta-device module with a custom param_init_fn before wrapping in FSDP, and how to use torchdistX's deferred_init for deferred initialization. The example includes a placeholder for module-type-specific initialization and prints the device of the first parameter.

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

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

### Get the name of an autograd node with Node.name()

Source: https://docs.pytorch.org/docs/2.14/generated/torch.autograd.graph.Node.name.html

Shows how to call Node.name() on a grad_fn instance. The example creates a tensor with requires_grad=True, clones it, and prints the name of the resulting grad_fn, which is 'CloneBackward0'.

```python
>>> import torch
>>> a = torch.tensor([0., 0., 0.], requires_grad=True)
>>> b = a.clone()
>>> assert isinstance(b.grad_fn, torch.autograd.graph.Node)
>>> print(b.grad_fn.name())
CloneBackward0
```

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

### Create QConfig with MinMaxObserver and default_observer

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

Example showing how to instantiate QConfig with observer classes, using with_args to set dtype to torch.qint8 for both activation and weight.

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

```

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

### Define forward with separate setup_context

Source: https://docs.pytorch.org/docs/2.14/generated/torch.autograd.function.InplaceFunction.html

Shows the second usage pattern for defining the forward method, where forward no longer accepts ctx and a separate setup_context staticmethod is required to set up the context.

```python
@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass


@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

```

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

### Example: Using storage_offset() on a tensor and a slice

Source: https://docs.pytorch.org/docs/2.14/generated/torch.Tensor.storage_offset.html

Shows that a full tensor has an offset of 0, while a slice starting at index 3 has an offset of 3. The output is displayed as Python interpreter results.

```python
>>> x = torch.tensor([1, 2, 3, 4, 5])
>>> x.storage_offset()
0
>>> x[3:].storage_offset()
3
```

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

### torch.nested.narrow example with jagged layout

Source: https://docs.pytorch.org/docs/2.14/nested.html

Narrows a nested tensor along dimension 1 using starts and lengths, with layout=torch.jagged. The resulting nested tensor is not contiguous, as shown by is_contiguous() returning False.

```python
>>> starts = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64)
>>> lengths = torch.tensor([3, 2, 2, 1, 5], dtype=torch.int64)
>>> narrow_base = torch.randn(5, 10, 20)
>>> nt_narrowed = torch.nested.narrow(narrow_base, 1, starts, lengths, layout=torch.jagged)
>>> nt_narrowed.is_contiguous()
False
```

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

### CUDAExtension

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

Creates a setuptools.Extension for CUDA/C++ with minimal arguments, including CUDA include path, library path, and runtime library. All arguments are forwarded to the setuptools.Extension constructor.

```APIDOC
## CUDAExtension

### Description
Creates a `setuptools.Extension` for CUDA/C++ with the bare minimum arguments to build a CUDA/C++ extension. This includes the CUDA include path, library path, and runtime library. All arguments are forwarded to the `setuptools.Extension` constructor.

### Method
Constructor

### Signature
`torch.utils.cpp_extension.CUDAExtension(name, sources, *args, **kwargs)`

### Parameters
- **name** (str) - Required - The name of the extension.
- **sources** (list[str]) - Required - List of source file paths.
- **\*args** - Optional - Additional positional arguments forwarded to `setuptools.Extension`.
- **\*\*kwargs** - Optional - Additional keyword arguments forwarded to `setuptools.Extension`.

### Example
```python
>>> from setuptools import setup
>>> from torch.utils.cpp_extension import BuildExtension, CUDAExtension
>>> setup(
...     name='cuda_extension',
...     ext_modules=[
...         CUDAExtension(
...                 name='cuda_extension',
...                 sources=['extension.cpp', 'extension_kernel.cu'],
...                 extra_compile_args={'cxx': ['-g'],
...                                     'nvcc': ['-O2']},
...                 extra_link_args=['-Wl,--no-as-needed', '-lcuda'])
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })
```

### Notes
- The PyTorch python API cannot be built with `py_limited_api=True`. If used, the user must avoid libtorch_python APIs and only use libtorch APIs.
- PyTorch defines `-DPy_LIMITED_API=min_supported_cpython` when `py_limited_api` is specified.
```

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

### Run decompositions with custom decomp_table

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

Shows how to get the default decompositions, override a specific operator with a custom decomposition, and then run decompositions on the exported program. The example uses torch.export.default_decompositions() and assigns your_custom_decomp to your_op.

```python
ep = torch.export.export(model, ...)
decomp_table = torch.export.default_decompositions()
decomp_table[your_op] = your_custom_decomp
ep = ep.run_decompositions(decomp_table=decomp_table)

```

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

### example_arguments()

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

Returns example arguments for the model's forward method. This method must be implemented by subclasses to provide sample inputs for tracing, testing, and ONNX export.

```APIDOC
## example_arguments()

### Description
Returns example arguments for the model's forward method. This method must be implemented by subclasses to provide sample inputs that can be used for tracing, testing, and ONNX export. The returned arguments should be representative of the expected input shapes and types during inference.

### Method
N/A (Python method)

### Endpoint
N/A

### Parameters
None

### Returns
- **Return type**: A tuple containing
- A tuple of positional arguments to pass to the forward method
- A dictionary of keyword arguments (or None if no kwargs are needed)

### Example
```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
    }
```
```

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

### SavePlanner._set_up_planner

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

Abstract method to initialize the planner to save a state_dict. Called on all ranks.

```APIDOC
## _set_up_planner(_state_dict_, _storage_meta=None_, _is_coordinator=False_)

### Description
Initialize this planner to save `state_dict`. Implementations should save those values as they won't be provided later in the save process. This is called on all ranks.

### Method
_abstract method_

### Signature
`_set_up_planner(_state_dict_, _storage_meta=None_, _is_coordinator=False_)`

### Parameters
- **state_dict** (dict) - The state dict to save.
- **storage_meta** (_StorageMeta_, optional) - Storage metadata.
- **is_coordinator** (bool, optional) - Whether this rank is the coordinator.

### Returns
None

```

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

### Correct CUDA stream usage with wait_stream and record_stream

Source: https://docs.pytorch.org/docs/2.14/notes/cuda.html

Fixed version of the previous example. Use s.wait_stream(torch.cuda.default_stream(cuda)) to ensure normal_() finishes before sum(A) starts, and A.record_stream(s) to prevent deallocation of A before sum(A) completes.

```python
cuda = torch.device('cuda')
s = torch.cuda.Stream()  # Create a new stream.
A = torch.empty((100, 100), device=cuda).normal_(0.0, 1.0)
s.wait_stream(torch.cuda.default_stream(cuda))  # NEW!
with torch.cuda.stream(s):
    B = torch.sum(A)
A.record_stream(s)  # NEW!
```

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

### List entrypoints with torch.hub.list()

Source: https://docs.pytorch.org/docs/2.14/hub.html

Lists all callable entrypoints in a GitHub repo. The example uses force_reload=True to discard cache and force a fresh download.

```python
>>> entrypoints = torch.hub.list("pytorch/vision", force_reload=True)

```
Tokens:
Raw