extension

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

Extension

This section provides insights into extending PyTorch’s capabilities. It covers custom operations, frontend APIs, and advanced topics like C++ extensions and dispatcher usage.

---

PyTorch Custom Operators Landing Page

This is the landing page for all things related to custom operators in PyTorch.

Extending-PyTorch,Frontend-APIs,C++,CUDA

Custom Python Operators

Create Custom Operators in Python. Useful for black-boxing a Python function for use with torch.compile.

Extending-PyTorch,Frontend-APIs,C++,CUDA

Custom C++ and CUDA Operators

How to extend PyTorch with custom C++ and CUDA operators.

Extending-PyTorch,Frontend-APIs,C++,CUDA

Custom Function Tutorial: Double Backward

Learn how to write a custom autograd Function that supports double backward.

Extending-PyTorch,Frontend-APIs

Custom Function Tutorial: Fusing Convolution and Batch Norm

Learn how to create a custom autograd Function that fuses batch norm into a convolution to improve memory usage.

Extending-PyTorch,Frontend-APIs

Custom C++ and CUDA Extensions

Create a neural network layer with no parameters using numpy. Then use scipy to create a neural network layer that has learnable weights.

Extending-PyTorch,Frontend-APIs,C++,CUDA

Registering a Dispatched Operator in C++

The dispatcher is an internal component of PyTorch which is responsible for figuring out what code should actually get run when you call a function like torch::add.

Extending-PyTorch,Frontend-APIs,C++

Extending Dispatcher For a New Backend in C++

Learn how to extend the dispatcher to add a new device living outside of the pytorch/pytorch repo and maintain it to keep in sync with native PyTorch devices.

Extending-PyTorch,Frontend-APIs,C++

Facilitating New Backend Integration by PrivateUse1

Learn how to integrate a new backend living outside of the pytorch/pytorch repo and maintain it to keep in sync with the native PyTorch backend.

Extending-PyTorch,Frontend-APIs,C++