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