torch.library

API регистрации операторов Python предоставляет возможности для расширения основного библиотеки операторов PyTorch с помощью пользовательских операторов. В настоящее время это можно сделать двумя способами:

1. Создание новых библиотек

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

     • Рассмотрим регистрацию нового оператора add в вашем недавно созданном пространстве имен foo. Вы можете получить доступ к этому оператору, используя API torch.ops и вызвав его, вызвав torch.ops.foo.add. Вы также можете получить доступ к определенным зарегистрированным перегрузкам, вызвав torch.ops.foo.add.{overload_name}.
     • Если вы зарегистрировали новое ядро для ключа диспетчеризации CUDA для этого оператора, то ваша настраиваемая функция будет вызвана для входных тензоров CUDA.
   • Это можно сделать, создав объекты класса Library типа "DEF".

2. Расширение существующих библиотек C++ (например, aten)

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

   • Это может быть полезно для заполнения пробелов в поддержке операторов для функции, реализованной с помощью ключа диспетчеризации. Например,

     • Вы можете добавить поддержку оператора для мета-тензоров (зарегистрировав функцию в ключе диспетчеризации Meta).
   • Это можно сделать, создав объекты класса Library типа "IMPL".

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

class torch.library.Library(ns, kind, dispatch_key='') [source]

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

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

define(schema, alias_analysis='') [source]

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

Параметры:

• schema – схема функции для определения нового оператора.
• alias_analysis (необязательно) – Указывает, можно ли вывести свойства алиасов аргументов оператора из схемы (поведение по умолчанию) или нет («CONSERVATIVE»).

Возвращает:

имя оператора, полученное из схемы.

Пример::

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

impl(op_name, fn, dispatch_key='') [source]

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

Параметры:

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

Пример::

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

Мы также добавили некоторые декораторы функций, чтобы упростить регистрацию функций для операторов:

• torch.library.impl()
• torch.library.define()