torch.library

Python operator registration API provides capabilities for extending PyTorch’s core library of operators with user defined operators. Currently, this can be done in two ways:

1. Creating new libraries

   • Lets you to register new operators and kernels for various backends and functionalities by specifying the appropriate dispatch keys. For example,

     • Consider registering a new operator add in your newly created namespace foo. You can access this operator using the torch.ops API and calling into by calling torch.ops.foo.add. You can also access specific registered overloads by calling torch.ops.foo.add.{overload_name}.

     • If you registered a new kernel for the CUDA dispatch key for this operator, then your custom defined function will be called for CUDA tensor inputs.

   • This can be done by creating Library class objects of "DEF" kind.

2. Extending existing C++ libraries (e.g., aten)

   • Lets you register kernels for existing operators corresponding to various backends and functionalities by specifying the appropriate dispatch keys.

   • This may come in handy to fill up spotty operator support for a feature implemented through a dispatch key. For example.,

     • You can add operator support for Meta Tensors (by registering function to the Meta dispatch key).

   • This can be done by creating Library class objects of "IMPL" kind.

A tutorial that walks you through some examples on how to use this API is available on Google Colab.

Warning

Dispatcher is a complicated PyTorch concept and having a sound understanding of Dispatcher is crucial to be able to do anything advanced with this API. This blog post is a good starting point to learn about Dispatcher.

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

A class to create libraries that can be used to register new operators or override operators in existing libraries from Python. A user can optionally pass in a dispatch keyname if they only want to register kernels corresponding to only one specific dispatch key.

Parameters

• ns – library name

• kind – “DEF”, “IMPL” (default: “IMPL”)

• dispatch_key – PyTorch dispatch key (default: “”)

define(schema, alias_analysis='')

Takes a schema to define a new operator. Also, optionally takes alias_analysis argument to indicate if the aliasing properties of the arguments can be inferred from the schema (default behavior) or not (“CONSERVATIVE”).

Returns the name of the operator as inferred from the schema.

We have also added some function decorators to make it convenient to register functions for operators:

• torch.library.impl()

• torch.library.define()