Class Library#
Defined inFile library.h
Class Documentation#
- classLibrary#
This object provides the API for defining operators and providing implementations at dispatch keys.
Typically, atorch::Library is not allocated directly; instead it is created by theTORCH_LIBRARY() orTORCH_LIBRARY_IMPL() macros.
Most methods ontorch::Library return a reference to itself, supporting method chaining.
// Examples:TORCH_LIBRARY(torchvision,m){// m is a torch::Librarym.def("roi_align",...);...}TORCH_LIBRARY_IMPL(aten,XLA,m){// m is a torch::Librarym.impl("add",...);...}
Public Functions
- ~Library()=default#
- inlineLibrary&def(c10::FunctionSchema&&s,conststd::vector<at::Tag>&tags={},_RegisterOrVerifyrv=_RegisterOrVerify::REGISTER)&#
Declare an operator with a schema, but don’t provide any implementations for it.
You’re expected to then provide implementations using theimpl() method. All template arguments are inferred.
// Example:TORCH_LIBRARY(myops,m){m.def("add(Tensor self, Tensor other) -> Tensor");}
- Parameters:
raw_schema – The schema of the operator to be defined. Typically, this is a
constchar*string literal, but any type accepted bytorch::schema() is accepted here.
- inlineLibrary&def(constchar*raw_schema,conststd::vector<at::Tag>&tags={},_RegisterOrVerifyrv=_RegisterOrVerify::REGISTER)&#
- inlineLibrary&set_python_module(constchar*pymodule,constchar*context="")#
Declares that for all operators that are subsequently def’ed, their fake impls may be found in the given Python module (pymodule).
This registers some help text that is used if the fake impl cannot be found.
Args:
pymodule: the python module
context: We may include this in the error message.
- inlineLibrary&impl_abstract_pystub(constchar*pymodule,constchar*context="")#
Deprecated; use set_python_module instead.
- template<typenameNameOrSchema,typenameFunc>
inlineLibrary&def(NameOrSchema&&raw_name_or_schema,Func&&raw_f,conststd::vector<at::Tag>&tags={})&# Define an operator for a schema and then register an implementation for it.
This is typically what you would use if you aren’t planning on making use of the dispatcher to structure your operator implementation. It’s roughly equivalent to callingdef() and thenimpl(), but if you omit the schema of the operator, we will infer it from the type of your C++ function. All template arguments are inferred.
// Example:TORCH_LIBRARY(myops,m){m.def("add",add_fn);}
- Parameters:
raw_name_or_schema – The schema of the operator to be defined, or just the name of the operator if the schema is to be inferred from
raw_f. Typically aconstchar*literal.raw_f – The C++ function that implements this operator. Any valid constructor oftorch::CppFunction is accepted here; typically you provide a function pointer or lambda.
- template<typenameName,typenameFunc>
inlineLibrary&impl(Namename,Func&&raw_f,_RegisterOrVerifyrv=_RegisterOrVerify::REGISTER)&# Register an implementation for an operator.
You may register multiple implementations for a single operator at different dispatch keys (seetorch::dispatch()). Implementations must have a corresponding declaration (fromdef()), otherwise they are invalid. If you plan to register multiple implementations, DO NOT provide a function implementation when youdef() the operator.
// Example:TORCH_LIBRARY_IMPL(myops,CUDA,m){m.impl("add",add_cuda);}
- Parameters:
name – The name of the operator to implement. Do NOT provide schema here.
raw_f – The C++ function that implements this operator. Any valid constructor oftorch::CppFunction is accepted here; typically you provide a function pointer or lambda.
- c10::OperatorName_resolve(constchar*name)const#
- inlineLibrary&def(detail::SelectiveStr<false>,conststd::vector<at::Tag>&tags[
[maybe_unused]]={})&#
- inlineLibrary&def(detail::SelectiveStr<true>raw_schema,conststd::vector<at::Tag>&tags={})&#
- template<typenameFunc>
inlineLibrary&def(detail::SelectiveStr<false>,Func&&,conststd::vector<at::Tag>&tags[[maybe_unused]]={})&#
- template<typenameFunc>
inlineLibrary&def(detail::SelectiveStr<true>raw_name_or_schema,Func&&raw_f,conststd::vector<at::Tag>&tags={})&#
- template<typenameFunc>
inlineLibrary&impl(detail::SelectiveStr<false>,Func&&)&#
- template<typenameDispatch,typenameFunc>
inlineLibrary&impl(detail::SelectiveStr<false>,Dispatch&&,Func&&)&#
- template<typenameFunc>
inlineLibrary&impl_UNBOXED(detail::SelectiveStr<false>,Func*)&#
- template<typenameFunc>
inlineLibrary&impl(detail::SelectiveStr<true>name,Func&&raw_f)&#
- template<typenameDispatch,typenameFunc>
inlineLibrary&impl(detail::SelectiveStr<true>name,Dispatch&&key,Func&&raw_f)&#
- template<typenameFunc>
inlineLibrary&impl_UNBOXED(detail::SelectiveStr<true>,Func*)&#
- template<typenameFunc>
inlineLibrary&fallback(Func&&raw_f)&# Register a fallback implementation for all operators which will be used if there is not a specific implementation for an operator available.
There MUST be a DispatchKey associated with a fallback; e.g., only call this fromTORCH_LIBRARY_IMPL() with namespace
_.// Example:TORCH_LIBRARY_IMPL(_,AutogradXLA,m){// If there is not a kernel explicitly registered// for AutogradXLA, fallthrough to the next// available kernelm.fallback(torch::CppFunction::makeFallthrough());}// See aten/src/ATen/core/dispatch/backend_fallback_test.cpp// for a full example of boxed fallback
- Parameters:
raw_f – The function that implements the fallback. Unboxed functions typically do not work as fallback functions, as fallback functions must work for every operator (even though they have varying type signatures). Typical arguments areCppFunction::makeFallthrough() or CppFunction::makeFromBoxedFunction()
- template<classCurClass>
inlinetorch::class_<CurClass>class_(detail::SelectiveStr<true>className)#
- template<classCurClass>
inlinedetail::ClassNotSelectedclass_(detail::SelectiveStr<false>className)#
- voidreset()#
- template<classCurClass>
inlineclass_<CurClass>class_(detail::SelectiveStr<true>className)#
Friends
- friendclassdetail::TorchLibraryInit
- ~Library()=default#
