Operator Registration

The library API provides macros and classes for registering custom operators with PyTorch's dispatcher.

Macros

TORCH_LIBRARY

{doxygendefine}

Example:

cpp

TORCH_LIBRARY(myops, m) {
  m.def("add(Tensor self, Tensor other) -> Tensor", &add_impl);
  m.def("mul(Tensor self, Tensor other) -> Tensor");
  m.impl("mul", torch::kCPU, &mul_cpu_impl);
  m.impl("mul", torch::kCUDA, &mul_cuda_impl);
}

TORCH_LIBRARY_IMPL

{doxygendefine}

Example:

cpp

TORCH_LIBRARY_IMPL(myops, XLA, m) {
  m.impl("mul", &mul_xla_impl);
}

TORCH_LIBRARY_FRAGMENT

{doxygendefine}

Example:

cpp

// In file1.cpp
TORCH_LIBRARY(myops, m) {
  m.def("add(Tensor self, Tensor other) -> Tensor", &add_impl);
}

// In file2.cpp
TORCH_LIBRARY_FRAGMENT(myops, m) {
  m.def("mul(Tensor self, Tensor other) -> Tensor", &mul_impl);
}

Classes

Library

{doxygenclass}

Example:

cpp

TORCH_LIBRARY(myops, m) {
  // Define with implementation
  m.def("add(Tensor self, Tensor other) -> Tensor", &add_impl);

  // Define schema only
  m.def("mul(Tensor self, Tensor other) -> Tensor");

  // Provide backend-specific implementations
  m.impl("mul", torch::kCPU, &mul_cpu_impl);
  m.impl("mul", torch::kCUDA, &mul_cuda_impl);
}

CppFunction

{doxygenclass}

:members:
:no-link:

OrderedDict

{doxygenclass}

:members:
:undoc-members:

Functions

The library API provides builder methods on the Library class for registering operators. See the Library class documentation above for the full API including def(), impl(), and fallback() methods.

Dispatch Keys

Common dispatch keys used with torch::dispatch():

torch::kCPU - CPU backend
torch::kCUDA - CUDA backend
torch::kAutograd - Autograd backend
torch::kMeta - Meta tensor backend