torch.library

Создано: Jun 13, 2022 | Последнее обновление: May 19, 2026

torch.library — это набор API для расширения основной библиотеки операторов PyTorch. Он содержит утилиты для тестирования пользовательских операторов, создания новых пользовательских операторов и расширения операторов, определённых с помощью API регистрации операторов C++ в PyTorch (например, операторов aten).

Подробное руководство по эффективному использованию этих API см. на странице PyTorch Custom Operators Landing Page, где приведены дополнительные сведения об их эффективном использовании.

Тестирование пользовательских операторов

Используйте torch.library.opcheck() для проверки пользовательских операторов на неправильное использование API Python torch.library и/или C++ TORCH_LIBRARY. Кроме того, если ваш оператор поддерживает обучение, используйте torch.autograd.gradcheck(), чтобы проверить математическую корректность градиентов.

torch.library.opcheck(op, args, kwargs=None, *, test_utils=('test_schema', 'test_autograd_registration', 'test_faketensor', 'test_aot_dispatch_dynamic'), raise_exception=True, atol=None, rtol=None) [исходный код]

Принимает оператор и несколько примеров аргументов и проверяет, правильно ли зарегистрирован оператор.

Иными словами, когда вы используете API torch.library/TORCH_LIBRARY для создания пользовательского оператора, вы задаёте метаданные о нём (например, сведения об изменяемости), а эти API требуют, чтобы передаваемые им функции удовлетворяли определённым свойствам (например, чтобы фиктивное/мета/абстрактное ядро не обращалось к указателю данных). opcheck проверяет эти метаданные и свойства.

В частности, проверяются следующие условия:

test_schema: соответствует ли схема реализации оператора. Например, если в схеме указано, что тензор изменяется, мы проверяем, изменяет ли его реализация. Если в схеме указано, что возвращается новый тензор, мы проверяем, возвращает ли реализация новый тензор (а не существующий тензор или его представление).
test_autograd_registration: если оператор поддерживает обучение (autograd), мы проверяем, зарегистрирована ли его формула autograd через torch.library.register_autograd или вручную для одного или нескольких ключей DispatchKey::Autograd. Любая другая регистрация на основе DispatchKey может привести к неопределённому поведению.
test_faketensor: есть ли у оператора ядро FakeTensor (и корректно ли оно). Ядро FakeTensor необходимо (но недостаточно) для работы оператора с API компиляции PyTorch (torch.compile/export/FX). Мы проверяем, зарегистрировано ли для оператора ядро FakeTensor (иногда также называемое метаядром) и корректно ли оно. В этом тесте сравниваются результаты выполнения оператора с обычными тензорами и FakeTensor: проверяется, совпадают ли метаданные тензоров (размеры/шаги/dtype/устройство и т. д.).
test_aot_dispatch_dynamic: корректно ли оператор работает с API компиляции PyTorch (torch.compile/export/FX). Проверяется, совпадают ли выходные данные (и градиенты, если применимо) при выполнении PyTorch в eager-режиме и torch.compile. Этот тест включает в себя test_faketensor и является сквозным тестом; кроме того, он проверяет поддержку функционализации оператором и поддержку FakeTensor и функционализации обратным проходом (если он существует).

Для наилучших результатов вызывайте opcheck несколько раз с репрезентативным набором входных данных. Если оператор поддерживает autograd, используйте opcheck с входными данными, имеющими requires_grad = True; если оператор поддерживает несколько устройств (например, CPU и CUDA), используйте opcheck с входными данными на всех поддерживаемых устройствах.

Параметры:

op (OpOverload | OpOverloadPacket | CustomOpDef) – Оператор. Должен быть либо функцией, декорированной с помощью torch.library.custom_op(), либо OpOverload/OpOverloadPacket из torch.ops.* (например, torch.ops.aten.sin, torch.ops.mylib.foo)
args (tuple[Any, ...]) – Позиционные аргументы оператора
kwargs (dict[str, Any] | None) – Именованные аргументы оператора
test_utils (str | Sequence[str]) – Тесты, которые следует запустить. По умолчанию запускаются все. Пример: (“test_schema”, “test_faketensor”)
raise_exception (bool) – Нужно ли выбрасывать исключение при первой ошибке. Если значение False, будет возвращён словарь с информацией о результатах каждого теста.
rtol (Optional[float]) – Относительный допуск при сравнении чисел с плавающей точкой. Если он задан, необходимо также задать atol. Если он не задан, выбираются значения по умолчанию на основе dtype (см. таблицу в torch.testing.assert_close()).
atol (Optional[float]) – Абсолютный допуск при сравнении чисел с плавающей точкой. Если он задан, необходимо также задать rtol. Если он не задан, выбираются значения по умолчанию на основе dtype (см. таблицу в torch.testing.assert_close()).

Тип возвращаемого значения:

dict[str, str]

Предупреждение

opcheck и torch.autograd.gradcheck() проверяют разные аспекты: opcheck проверяет корректность использования API torch.library, а torch.autograd.gradcheck() проверяет математическую корректность формулы autograd. Используйте оба теста для проверки пользовательских операторов, поддерживающих вычисление градиентов.

Пример

>>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=())
>>> def numpy_mul(x: Tensor, y: float) -> Tensor:
>>>     x_np = x.numpy(force=True)
>>>     z_np = x_np * y
>>>     return torch.from_numpy(z_np).to(x.device)
>>>
>>> @numpy_mul.register_fake
>>> def _(x, y):
>>>     return torch.empty_like(x)
>>>
>>> def setup_context(ctx, inputs, output):
>>>     y, = inputs
>>>     ctx.y = y
>>>
>>> def backward(ctx, grad):
>>>     return grad * ctx.y, None
>>>
>>> numpy_mul.register_autograd(backward, setup_context=setup_context)
>>>
>>> sample_inputs = [
>>>     (torch.randn(3), 3.14),
>>>     (torch.randn(2, 3, device='cuda'), 2.718),
>>>     (torch.randn(1, 10, requires_grad=True), 1.234),
>>>     (torch.randn(64, 64, device='cuda', requires_grad=True), 90.18),
>>> ]
>>>
>>> for args in sample_inputs:
>>>     torch.library.opcheck(numpy_mul, args)

Создание новых пользовательских операторов на Python

Для создания новых пользовательских операторов используйте torch.library.custom_op().

Выбор типа пользовательского оператора

При определении пользовательского оператора сначала решите, как он работает с алиасами и изменением данных. Этот выбор определяет mutates_args, tags, схему возвращаемого значения и поведение фиктивного ядра.

Does the operator mutate any Tensor input?
  |
  +-- no --> Functional operator
  |         Use mutates_args=().
  |         Return new Tensor objects that do not alias inputs or each other.
  |
  +-- yes --> Does it mutate the first positional Tensor argument and return it?
        |
        +-- yes --> In-place operator
        |          Use tags=torch.Tag.inplace.
        |          Use mutates_args with exactly the first argument name.
        |          Return that same first argument.
        |
        +-- no --> Does it mutate keyword-only Tensor out= arguments and return them?
              |
              +-- yes --> out= operator
              |          Use tags=torch.Tag.out.
              |          Do not read from the out= tensors.
              |          Put mutable output tensors after * as keyword-only args.
              |          Return all out= tensors in declaration order.
              |
              +-- no --> Mutable operator
                         Use mutates_args for the mutated arguments.
                         Do not return mutated inputs or aliases of inputs.

Функциональные операторы — самые простые, и при возможности следует отдавать предпочтение им. Они не должны изменять входные данные, а их результаты должны быть новыми значениями, не имеющими общих данных ни с одним входным значением или другим результатом.

Используйте tags=torch.Tag.inplace только для стандартной формы операции на месте: первый позиционный аргумент — это Tensor, он является единственным изменяемым аргументом, и оператор возвращает тот же Tensor. Например, схема пользовательского оператора на месте, выведенная автоматически, имеет следующий вид:

(Tensor(a!) x, ...) -> Tensor(a!)

Используйте tags=torch.Tag.out только для стандартных вариантов out=: все изменяемые выходные тензоры должны быть именованными аргументами, доступными только по ключевому слову, а оператор должен возвращать эти выходные данные в порядке их объявления. Реализация должна записывать данные в тензоры out=, но не должна считывать их. Например:

(Tensor x, *, Tensor(a!) out) -> Tensor(a!)
(Tensor x, *, Tensor(a!) out0, Tensor(b!) out1) -> (Tensor(a!), Tensor(b!))

Пользовательские операторы на месте и out= получают автоматически сгенерированное фиктивное ядро на основе их тега и схемы. Обычно регистрировать для таких операторов фиктивное ядро не требуется, однако это разрешено, если нужно переопределить автоматически сгенерированное поведение. Пользовательские операторы out= подчиняются тем же правилам autograd, что и встроенные операторы out=: они не поддерживают автоматическое дифференцирование, если для какого-либо аргумента требуется градиент.

Произвольное использование алиасов, например возврат алиаса входного значения, не относящегося к описанным выше шаблонам операций на месте или out=, не поддерживается в custom_op. Преобразования компилятора PyTorch, такие как функционализация, должны точно учитывать изменения данных и алиасы, поэтому использование алиасов в пользовательских операторах намеренно ограничено стандартными формами, описанными здесь. Если оператор изменяет входные данные, но не является оператором на месте или оператором out=, смоделируйте его как изменяемый оператор и не возвращайте алиасы входных данных.

Одна и та же операция может иметь разные контракты в зависимости от способа обработки выходных данных:

from torch import Tensor


@torch.library.custom_op(
    "mylib::add_inplace",
    mutates_args={"x"},
    tags=torch.Tag.inplace,
)
def add_inplace(x: Tensor, y: Tensor) -> Tensor:
    x.add_(y)
    return x


@torch.library.custom_op(
    "mylib::add_out",
    mutates_args={"out"},
    tags=torch.Tag.out,
)
def add_out(x: Tensor, y: Tensor, *, out: Tensor) -> Tensor:
    out.copy_(x + y)
    return out


@torch.library.custom_op("mylib::add_mutate", mutates_args={"x"})
def add_mutate(x: Tensor, y: Tensor) -> None:
    x.add_(y)

Подробнее об изменении данных, алиасах и преобразованиях

custom_op требует точного контракта на изменение данных и использование алиасов, поскольку PyTorch использует этот контракт в FakeTensor, autograd, функционализации и torch.compile.

Для функционального пользовательского оператора PyTorch предполагает, что оператор не изменяет входные данные, а возвращаемые тензоры являются новыми значениями. Такой тип операторов проще всего обрабатывать преобразованиям PyTorch.

Для пользовательского оператора на месте torch.Tag.inplace задаёт PyTorch более строгий и конкретный контракт: первый аргумент типа Tensor изменяется, а возвращаемый Tensor — тот же объект. Благодаря этому PyTorch может вывести фиктивное поведение из схемы, не требуя отдельного фиктивного ядра.

Для пользовательского оператора out= torch.Tag.out задаёт PyTorch соответствующий контракт варианта out=: выходные тензоры, доступные только по ключевому слову, изменяются и возвращаются в порядке, указанном в схеме. Тензоры out= используются только как выходные буферы: реализация не должна использовать их текущие значения в качестве входных данных. Как и встроенные операторы out=, эти пользовательские операторы не поддерживают автоматическое дифференцирование, если для какого-либо аргумента требуется градиент.

Для других изменяемых пользовательских операторов используйте mutates_args, чтобы указать изменяемые аргументы, но не возвращайте изменённые входные данные или их алиасы. Произвольные алиасы сложно обрабатывать преобразованиям: функционализация преобразует изменяющие данные программы в функциональные и должна точно знать, какие возвращаемые тензоры совместно используют хранилище с какими входными данными. Поэтому custom_op напрямую поддерживает стандартные шаблоны алиасов для операций на месте и out=, но не моделирует произвольные связи представлений или алиасов.

torch.library.custom_op(name, fn=None, /, *, mutates_args, device_types=None, schema=None, tags=None) [исходный код]

Оборачивает функцию в пользовательский оператор.

Пользовательский оператор может потребоваться, например, для следующих целей: - Обернуть стороннюю библиотеку или пользовательское ядро, чтобы оно работало с подсистемами PyTorch, например Autograd. - Запретить трассировке torch.compile/export/FX заглядывать внутрь функции.

Этот API используется как декоратор функции (см. примеры). Для предоставленной функции необходимо указать подсказки типов: они нужны для взаимодействия с различными подсистемами PyTorch.

Параметры:

name (str) – Имя пользовательского оператора в формате “{namespace}::{name}”, например “mylib::my_linear”. Имя используется как стабильный идентификатор оператора в подсистемах PyTorch (например, torch.export, графах FX). Во избежание конфликтов имён используйте в качестве пространства имён название проекта; например, все пользовательские операторы в pytorch/fbgemm используют пространство имён “fbgemm”.
mutates_args (Iterable[str] or "unknown") – Имена аргументов, которые изменяет функция. Это значение ДОЛЖНО быть точным, иначе поведение не определено. Если указано “unknown”, предполагается, что оператор изменяет все входные данные.
device_types (str | None | Sequence[str]) – Тип или типы устройств, для которых допустима функция. Если тип устройства не указан, функция используется как реализация по умолчанию для всех типов устройств. Примеры: “cpu”, “cuda”. При регистрации реализации для конкретного устройства для оператора, не принимающего тензоры, оператор должен иметь аргумент “device: torch.device”.
schema (str | None) – Строка схемы оператора. Если указано None (рекомендуется), схема оператора будет выведена из его аннотаций типов. Рекомендуется позволить системе вывести схему, если для этого нет особых причин. Пример: “(Tensor x, int y) -> (Tensor, Tensor)”.
tags (Tag | Sequence[Tag] | None) – один или несколько тегов, назначаемых оператору. Используйте torch.Tag.inplace для операторов, которые изменяют первый аргумент Tensor и возвращают его. Используйте torch.Tag.out для операторов, которые изменяют выходные тензоры, доступные только по ключевому слову, и возвращают их. Как и встроенные операторы out= в PyTorch, пользовательские операторы torch.Tag.out не поддерживают autograd.

Тип возвращаемого значения:

Callable[[Callable[[…], object]], CustomOpDef] | CustomOpDef

Для входных параметров оборачиваемой функции поддерживаются следующие типы:

Скалярные значения: int, float, bool, str, torch.types.Number
Тензоры: torch.Tensor
Перечисления/устройства: torch.dtype, torch.device
Плоский список элементов одного типа: list[torch.Tensor], list[int], list[float], list[bool], list[torch.types.Number]
Необязательные значения: Optional любого из перечисленных выше скалярных типов/типов тензоров
Типы, зарегистрированные через torch.library.register_custom_class()

Для возвращаемого значения поддерживаются следующие типы:

torch.Tensor, list[torch.Tensor], int, float, bool, torch.types.Number.

Примечание

Рекомендуется не передавать аргумент schema, а позволить системе вывести его из аннотаций типов. Самостоятельно задавать схему рискованно. Собственную схему можно указать, если интерпретация аннотации типа системой не соответствует вашим требованиям. Дополнительные сведения о записи строки схемы см. здесь

Примеры::

>>> import torch
>>> from torch import Tensor
>>> from torch.library import custom_op
>>> import numpy as np
>>>
>>> @custom_op("mylib::numpy_sin", mutates_args=())
>>> def numpy_sin(x: Tensor) -> Tensor:
>>>     x_np = x.cpu().numpy()
>>>     y_np = np.sin(x_np)
>>>     return torch.from_numpy(y_np).to(device=x.device)
>>>
>>> x = torch.randn(3)
>>> y = numpy_sin(x)
>>> assert torch.allclose(y, x.sin())
>>>
>>> # Example of a custom op that only works for one device type.
>>> @custom_op("mylib::numpy_sin_cpu", mutates_args=(), device_types="cpu")
>>> def numpy_sin_cpu(x: Tensor) -> Tensor:
>>>     x_np = x.numpy()
>>>     y_np = np.sin(x_np)
>>>     return torch.from_numpy(y_np)
>>>
>>> x = torch.randn(3)
>>> y = numpy_sin_cpu(x)
>>> assert torch.allclose(y, x.sin())
>>>
>>> # Example of a custom op that mutates an input
>>> @custom_op("mylib::numpy_sin_inplace", mutates_args={"x"}, device_types="cpu")
>>> def numpy_sin_inplace(x: Tensor) -> None:
>>>     x_np = x.numpy()
>>>     np.sin(x_np, out=x_np)
>>>
>>> x = torch.randn(3)
>>> expected = x.sin()
>>> numpy_sin_inplace(x)
>>> assert torch.allclose(x, expected)
>>>
>>> # Example of a custom op with inplace semantics
>>> @custom_op(
>>>     "mylib::numpy_sin_",
>>>     mutates_args={"x"},
>>>     device_types="cpu",
>>>     tags=torch.Tag.inplace,
>>> )
>>> def numpy_sin_(x: Tensor) -> Tensor:
>>>     x_np = x.numpy()
>>>     np.sin(x_np, out=x_np)
>>>     return x
>>>
>>> x = torch.randn(3)
>>> expected = x.sin()
>>> result = numpy_sin_(x)
>>> assert result is x
>>> assert torch.allclose(x, expected)
>>>
>>> # Example of a custom op with out= semantics
>>> @custom_op(
>>>     "mylib::numpy_sin_out",
>>>     mutates_args={"out"},
>>>     device_types="cpu",
>>>     tags=torch.Tag.out,
>>> )
>>> def numpy_sin_out(x: Tensor, *, out: Tensor) -> Tensor:
>>>     x_np = x.numpy()
>>>     out_np = out.numpy()
>>>     np.sin(x_np, out=out_np)
>>>     return out
>>>
>>> x = torch.randn(3)
>>> out = torch.empty_like(x)
>>> result = numpy_sin_out(x, out=out)
>>> assert result is out
>>> assert torch.allclose(out, x.sin())
>>>
>>> # Example of a factory function
>>> @torch.library.custom_op("mylib::bar", mutates_args={}, device_types="cpu")
>>> def bar(device: torch.device) -> Tensor:
>>>     return torch.ones(3)
>>>
>>> bar("cpu")
>>>
>>> # Example of a custom op with list inputs
>>> @custom_op("mylib::weighted_sum", mutates_args=())
>>> def weighted_sum(
>>>     tensors: list[Tensor],
>>>     weights: list[float],
>>> ) -> Tensor:
>>>     return sum(t * w for t, w in zip(tensors, weights))
>>>
>>> x = torch.randn(3)
>>> y = torch.randn(3)
>>> out = weighted_sum([x, y], [0.3, 0.7])

torch.library.triton_op(name, fn=None, /, *, mutates_args, schema=None) [исходный код]

Создаёт пользовательский оператор, реализация которого основана на одном или нескольких ядрах triton.

Это более структурированный способ использования ядер triton с PyTorch. Предпочтительно использовать ядра triton без обёрток пользовательских операторов torch.library (например, torch.library.custom_op(), torch.library.triton_op()), поскольку это проще. Используйте torch.library.custom_op()/torch.library.triton_op(), только если нужно создать оператор, который ведёт себя как встроенные операторы PyTorch. Например, можно использовать API-обёртку torch.library, чтобы определить поведение ядра triton при передаче ему подкласса тензора или при работе в TorchDispatchMode.

Если реализация состоит из одного или нескольких ядер triton, используйте torch.library.triton_op() вместо torch.library.custom_op(). torch.library.custom_op() рассматривает пользовательские операторы как непрозрачные объекты (torch.compile() и torch.export.export() никогда не трассируют их внутреннее содержимое), а triton_op делает реализацию видимой для этих подсистем, позволяя им оптимизировать ядра triton.

Обратите внимание: fn должен состоять только из вызовов операторов, распознаваемых PyTorch, и ядер triton. Любые ядра triton, вызываемые внутри fn, должны быть обёрнуты вызовом torch.library.wrap_triton().

Параметры:

name (str) – Имя пользовательского оператора в формате “{namespace}::{name}”, например “mylib::my_linear”. Имя используется как стабильный идентификатор оператора в подсистемах PyTorch (например, torch.export, графах FX). Во избежание конфликтов имён используйте в качестве пространства имён название проекта; например, все пользовательские операторы в pytorch/fbgemm используют пространство имён “fbgemm”.
mutates_args (Iterable[str] or "unknown") – Имена аргументов, которые изменяет функция. Это значение ДОЛЖНО быть точным, иначе поведение не определено. Если указано “unknown”, предполагается, что оператор изменяет все входные данные.
schema (str | None) – Строка схемы оператора. Если указано None (рекомендуется), схема оператора будет выведена из его аннотаций типов. Рекомендуется позволить системе вывести схему, если для этого нет особых причин. Пример: “(Tensor x, int y) -> (Tensor, Tensor)”.

Тип возвращаемого значения:

Callable

Пример:

>>> import torch
>>> from torch.library import triton_op, wrap_triton
>>>
>>> import triton
>>> from triton import language as tl
>>>
>>> @triton.jit
>>> def add_kernel(
>>>     in_ptr0,
>>>     in_ptr1,
>>>     out_ptr,
>>>     n_elements,
>>>     BLOCK_SIZE: "tl.constexpr",
>>> ):
>>>     pid = tl.program_id(axis=0)
>>>     block_start = pid * BLOCK_SIZE
>>>     offsets = block_start + tl.arange(0, BLOCK_SIZE)
>>>     mask = offsets < n_elements
>>>     x = tl.load(in_ptr0 + offsets, mask=mask)
>>>     y = tl.load(in_ptr1 + offsets, mask=mask)
>>>     output = x + y
>>>     tl.store(out_ptr + offsets, output, mask=mask)
>>>
>>> @triton_op("mylib::add", mutates_args={})
>>> def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
>>>     output = torch.empty_like(x)
>>>     n_elements = output.numel()
>>>
>>>     def grid(meta):
>>>         return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
>>>
>>>     # NB: we need to wrap the triton kernel in a call to wrap_triton
>>>     wrap_triton(add_kernel)[grid](x, y, output, n_elements, 16)
>>>     return output
>>>
>>> @torch.compile
>>> def f(x, y):
>>>     return add(x, y)
>>>
>>> x = torch.randn(3, device="cuda")
>>> y = torch.randn(3, device="cuda")
>>>
>>> z = f(x, y)
>>> assert torch.allclose(z, x + y)

torch.library.wrap_triton(triton_kernel, /) [исходный код]

Позволяет захватывать ядро triton в граф с помощью make_fx или non-strict torch.export.

Эти технологии выполняют трассировку на основе Dispatcher (через __torch_dispatch__) и не могут видеть вызовы необёрнутых ядер triton. API wrap_triton оборачивает ядро triton в вызываемый объект, который можно трассировать в граф.

Используйте этот API вместе с torch.library.triton_op().

Примеры

>>> import torch
>>> import triton
>>> from triton import language as tl
>>> from torch.fx.experimental.proxy_tensor import make_fx
>>> from torch.library import wrap_triton
>>>
>>> @triton.jit
>>> def add_kernel(
>>>     in_ptr0,
>>>     in_ptr1,
>>>     out_ptr,
>>>     n_elements,
>>>     BLOCK_SIZE: "tl.constexpr",
>>> ):
>>>     pid = tl.program_id(axis=0)
>>>     block_start = pid * BLOCK_SIZE
>>>     offsets = block_start + tl.arange(0, BLOCK_SIZE)
>>>     mask = offsets < n_elements
>>>     x = tl.load(in_ptr0 + offsets, mask=mask)
>>>     y = tl.load(in_ptr1 + offsets, mask=mask)
>>>     output = x + y
>>>     tl.store(out_ptr + offsets, output, mask=mask)
>>>
>>> def add(x, y):
>>>     output = torch.empty_like(x)
>>>     n_elements = output.numel()
>>>
>>>     def grid_fn(meta):
>>>         return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
>>>
>>>     wrap_triton(add_kernel)[grid_fn](x, y, output, n_elements, 16)
>>>     return output
>>>
>>> x = torch.randn(3, device="cuda")
>>> y = torch.randn(3, device="cuda")
>>> gm = make_fx(add)(x, y)
>>> print(gm.code)
>>> # def forward(self, x_1, y_1):
>>> #     empty_like = torch.ops.aten.empty_like.default(x_1, pin_memory = False)
>>> #     triton_kernel_wrapper_mutation_proxy = triton_kernel_wrapper_mutation(
>>> #         kernel_idx = 0, constant_args_idx = 0,
>>> #         grid = [(1, 1, 1)], kwargs = {
>>> #             'in_ptr0': x_1, 'in_ptr1': y_1, 'out_ptr': empty_like,
>>> #             'n_elements': 3, 'BLOCK_SIZE': 16
>>> #         })
>>> #     return empty_like

Тип возвращаемого значения:: Any

Расширение пользовательских операторов (созданных на Python или C++)

Используйте методы register.*, например torch.library.register_kernel() и torch.library.register_fake(), чтобы добавлять реализации любых операторов (они могли быть созданы с помощью torch.library.custom_op() или API регистрации операторов C++ в PyTorch).

torch.library.register_kernel(op, device_types, func=None, /, *, lib=None) [исходный код]

Зарегистрируйте реализацию этого оператора для типа устройства.

Допустимые значения device_types: «cpu», «cuda», «xla», «mps», «ipu», «xpu». Этот API можно использовать как декоратор.

Параметры:

op (str | OpOverload) – Оператор, для которого нужно зарегистрировать реализацию.
device_types (str | None | Sequence[str]) – Типы устройств, для которых нужно зарегистрировать реализацию. Если значение равно None, регистрация выполняется для всех типов устройств — используйте этот вариант, только если ваша реализация действительно не зависит от типа устройства.
func (Callable) – Функция, которую нужно зарегистрировать как реализацию для указанных типов устройств.
lib (Optional[Library]) – Если указано, срок действия этой регистрации

Примеры::

>>> import torch
>>> from torch import Tensor
>>> from torch.library import custom_op
>>> import numpy as np
>>>
>>> # Create a custom op that works on cpu
>>> @custom_op("mylib::numpy_sin", mutates_args=(), device_types="cpu")
>>> def numpy_sin(x: Tensor) -> Tensor:
>>>     x_np = x.numpy()
>>>     y_np = np.sin(x_np)
>>>     return torch.from_numpy(y_np)
>>>
>>> # Add implementations for the cuda device
>>> @torch.library.register_kernel("mylib::numpy_sin", "cuda")
>>> def _(x):
>>>     x_np = x.cpu().numpy()
>>>     y_np = np.sin(x_np)
>>>     return torch.from_numpy(y_np).to(device=x.device)
>>>
>>> x_cpu = torch.randn(3)
>>> x_cuda = x_cpu.cuda()
>>> assert torch.allclose(numpy_sin(x_cpu), x_cpu.sin())
>>> assert torch.allclose(numpy_sin(x_cuda), x_cuda.sin())

torch.library.register_autocast(op, device_type, cast_inputs, /, *, lib=None) [исходный код]

Зарегистрируйте правило диспетчеризации autocast для этого пользовательского оператора.

Допустимые значения device_type включают все типы устройств, поддерживающие autocast. Подробности см. в torch.amp.is_autocast_available().

Параметры:

op (str | OpOverload) – Оператор, для которого нужно зарегистрировать правило диспетчеризации autocast.
device_type (str) – Тип устройства. «cuda», «cpu», «xpu» или любой другой тип устройства, поддерживающий autocast. Тип совпадает со значением атрибута type объекта torch.device. Таким образом, тип устройства тензора можно получить с помощью Tensor.device.type.
cast_inputs (torch.dtype) – Когда пользовательский оператор выполняется в области с включённым autocast, преобразует входящие тензоры с плавающей точкой к целевому типу данных (тензоры, не являющиеся тензорами с плавающей точкой, не затрагиваются), а затем выполняет пользовательский оператор с отключённым autocast.
lib (Optional[Library]) – Если указано, срок действия этой регистрации

Примеры::

>>> import torch
>>> from torch import Tensor
>>> from torch.library import custom_op
>>>
>>> # Create a custom op that works on cuda
>>> @torch.library.custom_op("mylib::my_sin", mutates_args=())
>>> def my_sin(x: Tensor) -> Tensor:
>>>     return torch.sin(x)
>>>
>>> # Register autocast dispatch rule for the cuda device
>>> torch.library.register_autocast("mylib::my_sin", "cuda", torch.float16)
>>>
>>> x = torch.randn(3, dtype=torch.float32, device="cuda")
>>> with torch.autocast("cuda", dtype=torch.float16):
>>>     y = torch.ops.mylib.my_sin(x)
>>> assert y.dtype == torch.float16

torch.library.register_autograd(op, backward, /, *, setup_context=None, lib=None) [исходный код]

Зарегистрируйте формулу обратного прохода для этого пользовательского оператора.

Чтобы оператор работал с autograd, необходимо зарегистрировать формулу обратного прохода: 1. Нужно указать, как вычислять градиенты во время обратного прохода, предоставив функцию «backward». 2. Если для вычисления градиентов нужны какие-либо значения из прямого прохода, используйте setup_context для сохранения значений, необходимых обратному проходу.

backward выполняется во время обратного прохода. Она принимает (ctx, *grads): - grads — один или несколько градиентов. Количество градиентов совпадает с количеством выходных значений оператора. Объект ctx — это тот же объект ctx, который используется в torch.autograd.Function. Семантика backward_fn совпадает с семантикой torch.autograd.Function.backward().

setup_context(ctx, inputs, output) выполняется во время прямого прохода. Сохраняйте необходимые для обратного прохода значения в объекте ctx с помощью torch.autograd.function.FunctionCtx.save_for_backward() или присваивая их атрибутам ctx. Если у пользовательского оператора есть аргументы, доступные только по ключевому слову, ожидается, что сигнатура setup_context будет иметь вид setup_context(ctx, inputs, keyword_only_inputs, output).

И setup_context_fn, и backward_fn должны поддерживать трассировку. То есть они не должны напрямую обращаться к torch.Tensor.data_ptr() и не должны зависеть от глобального состояния или изменять его. Если вам нужна нетрассируемая функция обратного прохода, можно оформить её как отдельный custom_op и вызывать внутри backward_fn.

Если для разных устройств требуется разное поведение autograd, рекомендуем создать два разных пользовательских оператора — по одному для каждого устройства с отличающимся поведением — и выбирать нужный во время выполнения.

Примеры

>>> import torch
>>> import numpy as np
>>> from torch import Tensor
>>>
>>> @torch.library.custom_op("mylib::numpy_sin", mutates_args=())
>>> def numpy_sin(x: Tensor) -> Tensor:
>>>     x_np = x.cpu().numpy()
>>>     y_np = np.sin(x_np)
>>>     return torch.from_numpy(y_np).to(device=x.device)
>>>
>>> def setup_context(ctx, inputs, output) -> Tensor:
>>>     x, = inputs
>>>     ctx.save_for_backward(x)
>>>
>>> def backward(ctx, grad):
>>>     x, = ctx.saved_tensors
>>>     return grad * x.cos()
>>>
>>> torch.library.register_autograd(
...     "mylib::numpy_sin", backward, setup_context=setup_context
... )
>>>
>>> x = torch.randn(3, requires_grad=True)
>>> y = numpy_sin(x)
>>> (grad_x,) = torch.autograd.grad(y, x, torch.ones_like(y))
>>> assert torch.allclose(grad_x, x.cos())
>>>
>>> # Example with a keyword-only arg
>>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=())
>>> def numpy_mul(x: Tensor, *, val: float) -> Tensor:
>>>     x_np = x.cpu().numpy()
>>>     y_np = x_np * val
>>>     return torch.from_numpy(y_np).to(device=x.device)
>>>
>>> def setup_context(ctx, inputs, keyword_only_inputs, output) -> Tensor:
>>>     ctx.val = keyword_only_inputs["val"]
>>>
>>> def backward(ctx, grad):
>>>     return grad * ctx.val
>>>
>>> torch.library.register_autograd(
...     "mylib::numpy_mul", backward, setup_context=setup_context
... )
>>>
>>> x = torch.randn(3, requires_grad=True)
>>> y = numpy_mul(x, val=3.14)
>>> (grad_x,) = torch.autograd.grad(y, x, torch.ones_like(y))
>>> assert torch.allclose(grad_x, torch.full_like(x, 3.14))

torch.library.register_fake(op, func=None, /, *, lib=None, _stacklevel=1, allow_override=True) [исходный код]

Зарегистрируйте реализацию FakeTensor («фиктивную реализацию») для этого оператора.

Её также иногда называют «метаядром» или «абстрактной реализацией».

«Реализация FakeTensor» определяет поведение этого оператора для тензоров без данных («FakeTensor»). Для входных тензоров с определёнными свойствами (размерами/шагами/смещением в хранилище/устройством) она задаёт свойства выходных тензоров.

Реализация FakeTensor имеет ту же сигнатуру, что и оператор. Она выполняется как для FakeTensor, так и для метатензоров. При написании реализации FakeTensor считайте, что все тензорные входные данные оператора — это обычные тензоры CPU/CUDA/Meta, но у них нет хранилища, а результатом должны быть обычные тензоры CPU/CUDA/Meta. Реализация FakeTensor должна состоять только из операций PyTorch (и не должна напрямую обращаться к хранилищу или данным каких-либо входных или промежуточных тензоров).

Этот API можно использовать как декоратор (см. примеры).

Подробное руководство по пользовательским операторам см. на странице https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html

Параметры:

op_name – Имя оператора (вместе с вариантом перегрузки) или объект OpOverload.
func (Callable | None) – Реализация фиктивного тензора.
lib (Optional[Library]) – Библиотека, в которой нужно зарегистрировать фиктивный тензор.
allow_override (bool) – Флаг, определяющий, разрешено ли переопределить уже зарегистрированную фиктивную реализацию. По умолчанию переопределение разрешено; передайте False, чтобы вызвать ошибку, если для оператора уже есть фиктивная реализация. Это также применимо только в том случае, если пользовательский оператор создан не через torch.library.custom_op, так как переопределение существующей фиктивной реализации уже разрешено.

Примеры

>>> import torch
>>> import numpy as np
>>> from torch import Tensor
>>>
>>> # Example 1: an operator without data-dependent output shape
>>> @torch.library.custom_op("mylib::custom_linear", mutates_args=())
>>> def custom_linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor:
>>>     raise NotImplementedError("Implementation goes here")
>>>
>>> @torch.library.register_fake("mylib::custom_linear")
>>> def _(x, weight, bias):
>>>     assert x.dim() == 2
>>>     assert weight.dim() == 2
>>>     assert bias.dim() == 1
>>>     assert x.shape[1] == weight.shape[1]
>>>     assert weight.shape[0] == bias.shape[0]
>>>     assert x.device == weight.device
>>>
>>>     return (x @ weight.t()) + bias
>>>
>>> with torch._subclasses.fake_tensor.FakeTensorMode():
>>>     x = torch.randn(2, 3)
>>>     w = torch.randn(3, 3)
>>>     b = torch.randn(3)
>>>     y = torch.ops.mylib.custom_linear(x, w, b)
>>>
>>> assert y.shape == (2, 3)
>>>
>>> # Example 2: an operator with data-dependent output shape
>>> @torch.library.custom_op("mylib::custom_nonzero", mutates_args=())
>>> def custom_nonzero(x: Tensor) -> Tensor:
>>>     x_np = x.numpy(force=True)
>>>     res = np.stack(np.nonzero(x_np), axis=1)
>>>     return torch.tensor(res, device=x.device)
>>>
>>> @torch.library.register_fake("mylib::custom_nonzero")
>>> def _(x):
>>> # Number of nonzero-elements is data-dependent.
>>> # Since we cannot peek at the data in an fake impl,
>>> # we use the ctx object to construct a new symint that
>>> # represents the data-dependent size.
>>>     ctx = torch.library.get_ctx()
>>>     nnz = ctx.new_dynamic_size()
>>>     shape = [nnz, x.dim()]
>>>     result = x.new_empty(shape, dtype=torch.int64)
>>>     return result
>>>
>>> from torch.fx.experimental.proxy_tensor import make_fx
>>>
>>> x = torch.tensor([0, 1, 2, 3, 4, 0])
>>> trace = make_fx(torch.ops.mylib.custom_nonzero, tracing_mode="symbolic")(x)
>>> trace.print_readable()
>>>
>>> assert torch.allclose(trace(x), torch.ops.mylib.custom_nonzero(x))

torch.library.register_vmap(op, func=None, /, *, lib=None) [исходный код]

Зарегистрируйте реализацию vmap, чтобы поддержать torch.vmap() для этого пользовательского оператора.

Этот API можно использовать как декоратор (см. примеры).

Чтобы оператор работал с torch.vmap(), может потребоваться зарегистрировать реализацию vmap со следующей сигнатурой:

vmap_func(info, in_dims: Tuple[Optional[int]], *args, **kwargs),

где *args и **kwargs — аргументы и именованные аргументы для op. Тензорные аргументы, доступные только по ключевому слову, не поддерживаются.

Реализация определяет, как вычислить пакетную версию op для входных данных с дополнительным измерением (указанным параметром in_dims).

Для каждого аргумента в args у in_dims есть соответствующий Optional[int]. Его значение равно None, если аргумент не является тензором или над ним не выполняется vmap; в противном случае это целое число, указывающее измерение тензора, над которым выполняется vmap.

info — набор дополнительных метаданных, которые могут быть полезны: info.batch_size задаёт размер измерения, над которым выполняется vmap, а info.randomness — это параметр randomness, переданный в torch.vmap().

Функция func возвращает кортеж из (output, out_dims). Как и in_dims, out_dims должна иметь ту же структуру, что и output, и содержать по одному out_dim для каждого выходного значения, указывающему, содержит ли оно измерение vmap и на какой позиции оно находится.

Примеры

>>> import torch
>>> import numpy as np
>>> from torch import Tensor
>>> from typing import Tuple
>>>
>>> def to_numpy(tensor):
>>>     return tensor.cpu().numpy()
>>>
>>> lib = torch.library.Library("mylib", "FRAGMENT")
>>> @torch.library.custom_op("mylib::numpy_cube", mutates_args=())
>>> def numpy_cube(x: Tensor) -> Tuple[Tensor, Tensor]:
>>>     x_np = to_numpy(x)
>>>     dx = torch.tensor(3 * x_np ** 2, device=x.device)
>>>     return torch.tensor(x_np ** 3, device=x.device), dx
>>>
>>> def numpy_cube_vmap(info, in_dims, x):
>>>     result = numpy_cube(x)
>>>     return result, (in_dims[0], in_dims[0])
>>>
>>> torch.library.register_vmap(numpy_cube, numpy_cube_vmap)
>>>
>>> x = torch.randn(3)
>>> torch.vmap(numpy_cube)(x)
>>>
>>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=())
>>> def numpy_mul(x: Tensor, y: Tensor) -> Tensor:
>>>     return torch.tensor(to_numpy(x) * to_numpy(y), device=x.device)
>>>
>>> @torch.library.register_vmap("mylib::numpy_mul")
>>> def numpy_mul_vmap(info, in_dims, x, y):
>>>     x_bdim, y_bdim = in_dims
>>>     x = x.movedim(x_bdim, -1) if x_bdim is not None else x.unsqueeze(-1)
>>>     y = y.movedim(y_bdim, -1) if y_bdim is not None else y.unsqueeze(-1)
>>>     result = x * y
>>>     result = result.movedim(-1, 0)
>>>     return result, 0
>>>
>>>
>>> x = torch.randn(3)
>>> y = torch.randn(3)
>>> torch.vmap(numpy_mul)(x, y)

Примечание

Функция vmap должна сохранять семантику всего пользовательского оператора. То есть grad(vmap(op)) должна допускать замену на grad(map(op)).

Учитывайте это, если пользовательский оператор имеет особое поведение в обратном проходе.

torch.library.impl_abstract(qualname, func=None, *, lib=None, _stacklevel=1) [исходный код]: Этот API был переименован в torch.library.register_fake() в PyTorch 2.4. Вместо него используйте новое имя.

torch.library.get_ctx() [исходный код]

get_ctx() возвращает текущий объект AbstractImplCtx.

Вызов get_ctx() допустим только внутри фиктивной реализации (подробности использования см. в torch.library.register_fake().

Тип возвращаемого значения:: FakeImplCtx

torch.library.register_torch_dispatch(op, torch_dispatch_class, func=None, /, *, lib=None) [исходный код]

Регистрирует правило torch_dispatch для заданного оператора и torch_dispatch_class.

Это позволяет открыто регистрировать поведение оператора для torch_dispatch_class без необходимости изменять torch_dispatch_class или сам оператор.

torch_dispatch_class — это либо подкласс Tensor с __torch_dispatch__, либо TorchDispatchMode.

Если это подкласс Tensor, ожидается, что func будет иметь следующую сигнатуру: (cls, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any

Если это TorchDispatchMode, ожидается, что func будет иметь следующую сигнатуру: (mode, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any

args и kwargs будут нормализованы так же, как в __torch_dispatch__ (см. соглашение о вызове __torch_dispatch__).

Примеры

>>> import torch
>>>
>>> @torch.library.custom_op("mylib::foo", mutates_args={})
>>> def foo(x: torch.Tensor) -> torch.Tensor:
>>>     return x.clone()
>>>
>>> class MyMode(torch.utils._python_dispatch.TorchDispatchMode):
>>>     def __torch_dispatch__(self, func, types, args=(), kwargs=None):
>>>         return func(*args, **kwargs)
>>>
>>> @torch.library.register_torch_dispatch("mylib::foo", MyMode)
>>> def _(mode, func, types, args, kwargs):
>>>     x, = args
>>>     return x + 1
>>>
>>> x = torch.randn(3)
>>> y = foo(x)
>>> assert torch.allclose(y, x)
>>>
>>> with MyMode():
>>>     y = foo(x)
>>> assert torch.allclose(y, x + 1)

torch.library.infer_schema(prototype_function, /, *, mutates_args, op_name=None, tags=()) [исходный код]

Разбирает схему заданной функции с аннотациями типов. Схема выводится из аннотаций типов функции и может использоваться для определения нового оператора.

Мы делаем следующие предположения:

Ни один выходной объект не является псевдонимом входных объектов или других выходных объектов.

Вызывающий код (например, API пользовательских операторов) отвечает за проверку этих предположений.

Параметры:

prototype_function (Callable) – Функция, по аннотациям типов которой выводится схема.
op_name (Optional[str]) – Имя оператора в схеме. Если name равно None, имя не включается в выведенную схему. Обратите внимание, что входная схема для torch.library.Library.define требует указать имя оператора.
mutates_args ("unknown" | Iterable[str]) – Аргументы, изменяемые функцией.
tags (Tag | Sequence[Tag] | None) – одна или несколько меток для выведенной схемы. Используйте torch.Tag.inplace или torch.Tag.out, чтобы вывести стандартные псевдонимы для соответствующих типов операторов.

Возвращает:

Выведенная схема.

Тип возвращаемого значения:

str

Пример

>>> def foo_impl(x: torch.Tensor) -> torch.Tensor:
>>>     return x.sin()
>>>
>>> infer_schema(foo_impl, op_name="foo", mutates_args={})
foo(Tensor x) -> Tensor
>>>
>>> infer_schema(foo_impl, mutates_args={})
(Tensor x) -> Tensor

class torch._library.custom_ops.CustomOpDef(namespace, name, schema, fn, tags=None) [исходный код]

CustomOpDef — это обёртка над функцией, превращающая её в пользовательский оператор.

В ней есть различные методы для регистрации дополнительного поведения этого пользовательского оператора.

Не следует создавать экземпляр CustomOpDef напрямую; вместо этого используйте API torch.library.custom_op().

set_kernel_enabled(device_type, enabled=True) [исходный код]

Отключает или повторно включает уже зарегистрированное ядро для этого пользовательского оператора.

Если ядро уже отключено или включено, этот вызов ничего не делает.

Примечание

Если сначала отключить ядро, а затем зарегистрировать его, оно останется отключённым, пока его не включат снова.

Параметры:

device_type (str) – Тип устройства, для которого нужно отключить или включить ядро.
disable (bool) – Нужно ли отключить или включить ядро.

Пример

>>> inp = torch.randn(1)
>>>
>>> # define custom op `f`.
>>> @custom_op("mylib::f", mutates_args=())
>>> def f(x: Tensor) -> Tensor:
>>>     return torch.zeros(1)
>>>
>>> print(f(inp))  # tensor([0.]), default kernel
>>>
>>> @f.register_kernel("cpu")
>>> def _(x):
>>>     return torch.ones(1)
>>>
>>> print(f(inp))  # tensor([1.]), CPU kernel
>>>
>>> # temporarily disable the CPU kernel
>>> with f.set_kernel_enabled("cpu", enabled = False):
>>>     print(f(inp))  # tensor([0.]) with CPU kernel disabled

torch.library.get_kernel(op, dispatch_key) [исходный код]

Возвращает вычисленное ядро для заданного оператора и ключа диспетчеризации.

Эта функция извлекает ядро, которое будет выполнено для заданной комбинации оператора и ключа диспетчеризации. Возвращённый SafeKernelFunction можно использовать для вызова ядра в упакованном виде. Предполагается, что эта функция будет использоваться для получения исходного ядра для заданного ключа диспетчеризации, а затем для регистрации другого ядра с тем же ключом, которое в некоторых случаях вызывает исходное ядро.

Параметры:

op (str | OpOverload | CustomOpDef) – Имя оператора (вместе с вариантом перегрузки) или объект OpOverload. Может быть строкой (например, «aten::add.Tensor»), объектом OpOverload или CustomOpDef.
dispatch_key (str | torch.DispatchKey) – Ключ диспетчеризации, для которого нужно получить ядро. Может быть строкой (например, «CPU», «CUDA») или значением перечисления DispatchKey.

Возвращает:

Безопасную функцию ядра, которую можно использовать для: вызова ядра.

Тип возвращаемого значения:

torch._C._SafeKernelFunction

Вызывает исключение:

RuntimeError – Если оператор не существует.

Пример

>>> # Get the CPU kernel for torch.add
>>> kernel = torch.library.get_kernel("aten::add.Tensor", "CPU")
>>>
>>> # You can also use DispatchKey enum
>>> kernel = torch.library.get_kernel("aten::add.Tensor", torch.DispatchKey.CPU)
>>>
>>> # Or use an OpOverload directly
>>> kernel = torch.library.get_kernel(torch.ops.aten.add.Tensor, "CPU")
>>>
>>> # Example: Using get_kernel in a custom op with conditional dispatch
>>> # Get the original kernel for torch.sin
>>> original_sin_kernel = torch.library.get_kernel("aten::sin", "CPU")
>>>
>>> # If input has negative values, use original sin, otherwise return zeros
>>> def conditional_sin_impl(dispatch_keys, x):
>>>     if (x < 0).any():
>>>         return original_sin_kernel.call_boxed(dispatch_keys, x)
>>>     else:
>>>         return torch.zeros_like(x)
>>>
>>> lib = torch.library.Library("aten", "IMPL")
>>> # with_keyset=True so the first argument to the impl is the current DispatchKeySet
>>> which needs to be the first argument to ``kernel.call_boxed``
>>> lib.impl("sin", conditional_sin_impl, "CPU", with_keyset=True)
>>>
>>> # Test the conditional behavior
>>> x_positive = torch.tensor([1.0, 2.0])
>>> x_mixed = torch.tensor([-1.0, 2.0])
>>> torch.sin(x_positive)
tensor([0., 0.])
>>> torch.sin(x_mixed)
tensor([-0.8415, 0.9093])

Низкоуровневые API

Следующие API представляют собой прямые привязки к низкоуровневым API PyTorch на C++ для регистрации операторов.

Предупреждение

Низкоуровневые API регистрации операторов и диспетчер PyTorch — сложная концепция PyTorch. По возможности рекомендуем использовать приведённые выше API более высокого уровня (которым не требуется объект torch.library.Library). Эта запись в блоге — хорошая отправная точка для знакомства с диспетчером PyTorch.

Учебное пособие с примерами использования этого API доступно на Google Colab.

class torch.library.Library(ns, kind, dispatch_key='') [исходный код]

Класс для создания библиотек, с помощью которых из Python можно регистрировать новые операторы или переопределять операторы в существующих библиотеках. При необходимости пользователь может передать имя ключа диспетчеризации, если требуется регистрировать ядра только для одного конкретного ключа диспетчеризации.

Чтобы создать библиотеку для переопределения операторов в существующей библиотеке (с именем ns), укажите kind «IMPL». Чтобы создать новую библиотеку (с именем ns) для регистрации новых операторов, укажите kind «DEF». Чтобы создать фрагмент возможно существующей библиотеки для регистрации операторов (и обойти ограничение, согласно которому для данного пространства имён может существовать только одна библиотека), укажите kind «FRAGMENT».

Параметры:

ns – имя библиотеки
kind – «DEF», «IMPL», «FRAGMENT»
dispatch_key – ключ диспетчеризации PyTorch (по умолчанию: «»)

define(schema, alias_analysis='', *, tags=()) [исходный код]

Определяет новый оператор и его семантику в пространстве имён ns.

Параметры:

schema – схема функции для определения нового оператора.
alias_analysis (необязательно) – указывает, можно ли определить свойства алиасинга аргументов оператора по схеме (поведение по умолчанию) или нельзя («CONSERVATIVE»).
tags (Tag | Sequence[Tag]) – один или несколько torch.Tag, применяемых к этому оператору. Присвоение оператору тега меняет его поведение в различных подсистемах PyTorch; внимательно ознакомьтесь с документацией по torch.Tag, прежде чем применять тег.

Возвращает:

имя оператора, определённое по схеме.

Пример:

>>> my_lib = Library("mylib", "DEF")
>>> my_lib.define("sum(Tensor self) -> Tensor")

fallback(fn, dispatch_key='', *, with_keyset=False) [исходный код]

Регистрирует реализацию функции в качестве резервной для указанного ключа.

Эта функция работает только для библиотеки с глобальным пространством имён («_»).

Параметры:

fn – функция, используемая в качестве резервной для указанного ключа диспетчеризации, или fallthrough_kernel() для регистрации пропуска.
dispatch_key – ключ диспетчеризации, для которого следует зарегистрировать входную функцию. По умолчанию используется ключ диспетчеризации, с которым была создана библиотека.
with_keyset – флаг, определяющий, следует ли передавать набор ключей текущего вызова диспетчера в качестве первого аргумента функции fn. Это нужно для создания подходящего набора ключей при повторных вызовах диспетчеризации.

Пример:

>>> my_lib = Library("_", "IMPL")
>>> def fallback_kernel(op, *args, **kwargs):
>>>     # Handle all autocast ops generically
>>>     # ...
>>> my_lib.fallback(fallback_kernel, "Autocast")

impl(op_name, fn, dispatch_key='', *, with_keyset=False, allow_override=True) [исходный код]

Регистрирует реализацию функции для оператора, определённого в библиотеке.

Параметры:

op_name – имя оператора (вместе с вариантом перегрузки) или объект OpOverload.
fn – функция, реализующая оператор для входного ключа диспетчеризации, или fallthrough_kernel() для регистрации пропуска.
dispatch_key – ключ диспетчеризации, для которого следует зарегистрировать входную функцию. По умолчанию используется ключ диспетчеризации, с которым была создана библиотека.
with_keyset – флаг, определяющий, следует ли передавать набор ключей текущего вызова диспетчера в качестве первого аргумента функции fn. Это нужно для создания подходящего набора ключей при повторных вызовах диспетчеризации.
allow_override – флаг, определяющий, следует ли переопределить существующую зарегистрированную реализацию ядра. По умолчанию включён; передайте False, чтобы выдать ошибку, если для этого ключа диспетчеризации уже зарегистрировано ядро.

Пример::

>>> my_lib = Library("aten", "IMPL")
>>> def div_cpu(self, other):
>>>     return self * (1 / other)
>>> my_lib.impl("div.Tensor", div_cpu, "CPU")

register_symm_mem_args(op_name, arg_names) [исходный код]

Регистрирует аргументы, для которых требуется выделение симметричной памяти при вызове оператора.

Этот метод позволяет декларативно указать, какие аргументы оператора требуют обработки симметричной памяти. При использовании с torch.compile Inductor автоматически выделяет память для этих аргументов в области памяти NVLink, доступной через P2P (empty_strided_p2p), обеспечивая межграфический процессорный обмен данными без копирования.

Схема оператора должна содержать аргумент group_name (str). FallbackKernel._maybe_realize_symm_mem_args Inductor извлекает group_name во время компиляции, чтобы выделить буферы P2P для нужной группы процессов. Без group_name автоматическая материализация пропускается.

Параметры:

op_name – имя оператора (вместе с вариантом перегрузки) или объект OpOverload.
arg_names – список имён аргументов, для которых требуется выделение симметричной памяти.

Пример::

>>> my_lib = Library("symm_mem", "FRAGMENT")
>>> my_lib.define("one_shot_all_reduce(Tensor input, str reduce_op, str group_name) -> Tensor")
>>> my_lib.register_symm_mem_args("one_shot_all_reduce", ["input"])

torch.library.fallthrough_kernel() [исходный код]: Функция-заглушка, передаваемая в Library.impl для регистрации пропуска.

torch.library.define(qualname, schema, *, lib=None, tags=()) [исходный код]

torch.library.define(lib, schema, alias_analysis='')

Определяет новый оператор.

В PyTorch определение оператора (сокращённо «op») состоит из двух этапов: - необходимо определить оператор (указав его имя и схему); - необходимо реализовать поведение оператора при взаимодействии с различными подсистемами PyTorch, такими как тензоры CPU/CUDA, Autograd и т. д.

Эта точка входа определяет пользовательский оператор (первый этап); затем необходимо выполнить второй этап, вызвав различные API impl_*, например torch.library.impl() или torch.library.register_fake().

Параметры:

qualname (str) – полное имя оператора. Должно быть строкой вида «namespace::name», например «aten::sin». Операторам в PyTorch необходимо пространство имён, чтобы избежать совпадений имён; каждый оператор можно создать только один раз. Если вы разрабатываете библиотеку Python, рекомендуем использовать в качестве пространства имён имя модуля верхнего уровня.
schema (str) – схема оператора. Например, «(Tensor x) -> Tensor» для оператора, который принимает один Tensor и возвращает один Tensor. Схема не содержит имя оператора (оно передаётся в qualname).
lib (Optional[Library]) – если указан, время существования этого оператора будет ограничено временем существования объекта Library.
tags (Tag | Sequence[Tag]) – один или несколько torch.Tag, применяемых к этому оператору. Присвоение оператору тега меняет его поведение в различных подсистемах PyTorch; внимательно ознакомьтесь с документацией по torch.Tag, прежде чем применять тег.

Пример::

>>> import torch
>>> import numpy as np
>>>
>>> # Define the operator
>>> torch.library.define("mylib::sin", "(Tensor x) -> Tensor")
>>>
>>> # Add implementations for the operator
>>> @torch.library.impl("mylib::sin", "cpu")
>>> def f(x):
>>>     return torch.from_numpy(np.sin(x.numpy()))
>>>
>>> # Call the new operator from torch.ops.
>>> x = torch.randn(3)
>>> y = torch.ops.mylib.sin(x)
>>> assert torch.allclose(y, x.sin())

torch.library.impl(lib, name, dispatch_key='') [исходный код]

torch.library.impl(qualname:str, types:str|Sequence[str], func:None=None, *, lib:Library|None=None) → Callable[[Callable[...,object]],None]

torch.library.impl(qualname:str, types:str|Sequence[str], func:Callable[...,object], *, lib:Library|None=None) → None

torch.library.impl(lib:Library, name:str, dispatch_key:str='') → Callable[[Callable[_P,_T]],Callable[_P,_T]]

Регистрирует реализацию для типа устройства данного оператора.

Чтобы зарегистрировать эту реализацию в качестве реализации по умолчанию для ВСЕХ типов устройств, для types можно передать «default». Используйте это только в том случае, если реализация действительно поддерживает все типы устройств; например, это верно, если она представляет собой композицию встроенных операторов PyTorch.

Этот API можно использовать в качестве декоратора. С ним можно использовать вложенные декораторы при условии, что они возвращают функцию и размещены внутри этого API (см. пример 2).

Допустимые типы: «cpu», «cuda», «xla», «mps», «ipu», «xpu».

Параметры:

qualname (str) – должна быть строкой вида «namespace::operator_name».
types (str | Sequence[str]) – типы устройств, для которых регистрируется реализация.
lib (Optional[Library]) – если указан, время существования этой регистрации будет ограничено временем существования объекта Library.

Примеры

>>> import torch
>>> import numpy as np
>>> # Example 1: Register function.
>>> # Define the operator
>>> torch.library.define("mylib::mysin", "(Tensor x) -> Tensor")
>>>
>>> # Add implementations for the cpu device
>>> @torch.library.impl("mylib::mysin", "cpu")
>>> def f(x):
>>>     return torch.from_numpy(np.sin(x.numpy()))
>>>
>>> x = torch.randn(3)
>>> y = torch.ops.mylib.mysin(x)
>>> assert torch.allclose(y, x.sin())
>>>
>>> # Example 2: Register function with decorator.
>>> def custom_decorator(func):
>>>     def wrapper(*args, **kwargs):
>>>         return func(*args, **kwargs) + 1
>>>     return wrapper
>>>
>>> # Define the operator
>>> torch.library.define("mylib::sin_plus_one", "(Tensor x) -> Tensor")
>>>
>>> # Add implementations for the operator
>>> @torch.library.impl("mylib::sin_plus_one", "cpu")
>>> @custom_decorator
>>> def f(x):
>>>     return torch.from_numpy(np.sin(x.numpy()))
>>>
>>> # Call the new operator from torch.ops.
>>> x = torch.randn(3)
>>>
>>> y1 = torch.ops.mylib.sin_plus_one(x)
>>> y2 = torch.sin(x) + 1
>>> assert torch.allclose(y1, y2)

© 2026, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://docs.pytorch.org/docs/2.14/library.html