Open-Source Wikis

/

PyTorch

/

Primitives

/

Operators and schemas

pytorch/pytorch

Operators and schemas

What an operator is

An operator is a registered, dispatch-key-routed callable: e.g., aten::add, aten::matmul, c10d::all_reduce, my_lib::my_op. Each operator has:

  • A schema — a typed signature like aten::add(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor.
  • A set of kernels, one per dispatch key that has been registered.
  • An operator handle that the dispatcher uses to look it up.

Schema syntax

Schemas are strings with a custom mini-grammar parsed by c10::FunctionSchema (aten/src/ATen/core/function_schema.h). The grammar supports:

  • Argument types: Tensor, Tensor? (optional), Tensor[] (list), int, float, Scalar, bool, str, Device, Layout, MemoryFormat, Generator?, SymInt, SymFloat, SymBool, …
  • Argument modifiers: (a!) mutable alias, (a) immutable alias (view), default values.
  • * separator: arguments after * are keyword-only.
  • Return types: single, tuple, or named tuple (-> (Tensor sums, Tensor counts)).
  • Out-variant suffix: .out declares an explicit output argument.

Examples from native_functions.yaml:

aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
aten::add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> Tensor(a!)
aten::add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)
aten::cumsum(Tensor self, int dim, *, ScalarType? dtype=None) -> Tensor
aten::sum.dim_IntList(Tensor self, int[1] dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor

The names like add.Tensor, add.out, add.Scalar are overload names. They allow multiple ops with the same base name and different argument types to coexist.

Aliasing annotations

(a!) means "this argument may be mutated and aliases identifier a". (a) means "this argument may be aliased by the output but is not mutated". Used by autograd, functionalization, and the dispatcher's view-tracking logic to know which inputs/outputs share storage.

The Tensor(a!) syntax in add_ and add.out is what tells PyTorch these ops mutate input/output and how the alias relationship works.

Tags

Each schema entry can have tags (in aten/src/ATen/native/tags.yaml) that mark behavioural properties: pointwise, inplace_view, data_dependent_output, nondeterministic_seeded, view_copy, core, canonical, etc. Used by:

  • The compiler to know which ops can be reordered/fused.
  • Tests to know which OpInfo categories to apply.
  • Functionalization to know how to rewrite views.

Registration

There are three registration paths:

YAML-driven (most ops)

aten/src/ATen/native/native_functions.yaml is the canonical source. torchgen reads each entry and produces:

  • A C++ free function at::add(...) declaration in Operators.h.
  • A registration block in RegisterCPU.cpp, RegisterCUDA.cpp, etc. for each dispatch: row.
  • An autograd kernel registration if a derivatives.yaml entry exists.
  • A Python binding entry (python_torch_functions.cpp).

This is the path 99% of in-tree ops take.

TORCH_LIBRARY macros

TORCH_LIBRARY(my_lib, m) {
  m.def("my_op(Tensor x, int y) -> Tensor");
}
TORCH_LIBRARY_IMPL(my_lib, CPU, m) {
  m.impl("my_op", &my_op_cpu);
}
TORCH_LIBRARY_IMPL(my_lib, AutogradCPU, m) {
  m.impl("my_op", autograd::my_op_autograd);
}

Used for out-of-tree ops, JIT registrations, and dispatcher-key-specific kernels that don't fit the YAML pattern.

torch.library (Python)

import torch
@torch.library.custom_op("my_lib::my_op", mutates_args=())
def my_op(x: torch.Tensor) -> torch.Tensor:
    return x.relu()

Wraps TORCH_LIBRARY_IMPL from Python, with helpers for fake kernels, autograd formulas, and tag setting.

OperatorHandle and OperatorEntry

Internally, each operator is represented by an OperatorHandle (a stable handle returned by c10::Dispatcher::findOp) backing onto an OperatorEntry (aten/src/ATen/core/dispatch/OperatorEntry.h). The OperatorEntry holds:

  • The schema.
  • An array indexed by DispatchKey of registered kernels.
  • Backpointer to the dispatcher.
  • Cached fallthrough information.

When a kernel is registered for a key, the entry's array is updated; when the op is called the dispatcher walks the entry's array.

Boxed vs. unboxed call

Each kernel can be called via two paths:

  • Unboxed (typed): kernel.callUnboxed<ReturnT, Args...>(args...). The fast path used by the codegen-generated wrappers.
  • Boxed (IValue-based): kernel.callBoxed(stack). Slower; used by JIT, by torch.library Python kernels, and by introspection tools.

Most leaf kernels expose both via type erasure helpers in aten/src/ATen/core/boxing/.

Where to look

File Purpose
aten/src/ATen/native/native_functions.yaml The op registry
aten/src/ATen/native/README.md Op-author guide
aten/src/ATen/core/function_schema.h FunctionSchema
aten/src/ATen/core/op_registration/op_registration.h Programmatic registration
aten/src/ATen/core/dispatch/OperatorEntry.h Per-op kernel table
aten/src/ATen/core/Library.h TORCH_LIBRARY macros
aten/src/ATen/native/tags.yaml Op tags
tools/autograd/derivatives.yaml Differentiation rules

Built by Factory AutoWiki from public repository content. It is a generated preview for codebase exploration, not source-maintained documentation.

Operators and schemas – PyTorch wiki | Factory