Program Listing for File library.h¶
↰ Return to documentation for file (torch/library.h
)
// m is a torch::Library; methods on it will define
// operators in the myops namespace
m.def("add", add_impl);
}
m.fallback(xla_fallback);
}
// m is a torch::Library; methods on it will define
// CPU implementations of operators in the myops namespace.
// It is NOT valid to call torch::Library::def()
// in this context.
m.impl("add", add_cpu_impl);
}
// You must define all of the operators for this library in
// this namespace.
TORCH_LIBRARY(myops, m) {
// Define a operator with exactly one implementation for all backends.
m.def("add(Tensor self, Tensor other) -> Tensor", &add_impl);
// Define a schema for an operator, but provide no implementation
// (use this syntax if you want to use the dispatcher)
m.def("mul(Tensor self, Tensor other) -> Tensor");
// Provide an implementation for a defined operator (you can
// provide multiple; one per backend). The dispatcher takes care of
// calling the correct implementation depending on if we get a CPU
// tensor or a CUDA tensor
m.impl("mul", torch::kCPU, &mul_cpu_impl);
m.impl("mul", torch::kCUDA, &mul_cuda_impl);
}
// Define implementations for operators for a non-standard backend,
// e.g., XLA (valid values are entries of DispatchKey). This can
// be used to define operators in a different file than the initial
// TORCH_LIBRARY definition (e.g., if it is in an external library)
TORCH_LIBRARY_IMPL(myops, XLA, m) {
m.impl("mul", &mul_xla_impl);
}
#pragma once
#include <ATen/core/op_registration/infer_schema.h>
#include <ATen/core/op_registration/op_allowlist.h>
#include <ATen/core/dispatch/Dispatcher.h>
#include <c10/core/DispatchKey.h>
#include <torch/csrc/jit/frontend/function_schema_parser.h>
// Just for inferFunctionSchemaFromFunctor
#include <ATen/core/enum_tag.h>
#include <ATen/core/op_registration/op_registration.h>
namespace torch {
#if defined C10_MOBILE
struct NoInferSchemaTag {};
#endif
#define HAS_PT2_COMPLIANT_TAG
// For multipy/torchdeploy use case
enum class _RegisterOrVerify { REGISTER, VERIFY };
template <class CurClass>
class class_;
#define HAS_IMPL_ABSTRACT_PYSTUB
class TORCH_API CppFunction final {
// TODO: This is morally the same thing as KernelRegistrationConfig, but it's
// opaque to the user.
public:
template <typename Func>
explicit CppFunction(
Func* f,
std::enable_if_t<
c10::guts::is_function_type<Func>::value,
std::nullptr_t> = nullptr)
: func_(c10::KernelFunction::makeFromUnboxedRuntimeFunction(f)),
cpp_signature_(c10::impl::CppSignature::make<Func>()),
schema_(
c10::detail::inferFunctionSchemaFromFunctor<std::decay_t<Func>>()),
debug_() {}
template <typename FuncPtr>
explicit CppFunction(
FuncPtr f,
std::enable_if_t<
c10::is_compile_time_function_pointer<FuncPtr>::value,
std::nullptr_t> = nullptr)
: func_(c10::KernelFunction::makeFromUnboxedFunction(f)),
cpp_signature_(
c10::impl::CppSignature::make<typename FuncPtr::FuncType>()),
schema_(c10::detail::inferFunctionSchemaFromFunctor<
typename FuncPtr::FuncType>()),
debug_() {}
template <typename Lambda>
explicit CppFunction(
Lambda&& f,
std::enable_if_t<
c10::guts::is_functor<std::decay_t<Lambda>>::value,
std::nullptr_t> = nullptr)
: func_(c10::KernelFunction::makeFromUnboxedLambda(
std::forward<Lambda>(f))),
cpp_signature_(c10::impl::CppSignature::make<Lambda>()),
schema_(c10::detail::inferFunctionSchemaFromFunctor<
std::decay_t<Lambda>>()),
debug_() {}
#if defined C10_MOBILE
template <typename Func>
explicit CppFunction(
Func* f,
NoInferSchemaTag,
std::enable_if_t<
c10::guts::is_function_type<Func>::value,
std::nullptr_t> = nullptr)
: func_(c10::KernelFunction::makeFromUnboxedRuntimeFunction(f)),
cpp_signature_(c10::impl::CppSignature::make<Func>())
// TODO: Don't go through WrapRuntimeKernelFunctor
,
schema_(nullptr),
debug_() {}
template <typename FuncPtr>
explicit CppFunction(
FuncPtr f,
NoInferSchemaTag,
std::enable_if_t<
c10::is_compile_time_function_pointer<FuncPtr>::value,
std::nullptr_t> = nullptr)
: func_(c10::KernelFunction::makeFromUnboxedFunction(f)),
cpp_signature_(
c10::impl::CppSignature::make<typename FuncPtr::FuncType>())
// TODO: Don't go through WrapRuntimeKernelFunctor
,
schema_(nullptr),
debug_() {}
template <typename Lambda>
explicit CppFunction(
Lambda&& f,
NoInferSchemaTag,
std::enable_if_t<
c10::guts::is_functor<std::decay_t<Lambda>>::value,
std::nullptr_t> = nullptr)
: func_(c10::KernelFunction::makeFromUnboxedLambda(
std::forward<Lambda>(f))),
cpp_signature_(c10::impl::CppSignature::make<Lambda>())
// TODO: Don't go through WrapRuntimeKernelFunctor
,
schema_(nullptr),
debug_() {}
#endif
~CppFunction();
CppFunction(const CppFunction&) = delete;
CppFunction& operator=(const CppFunction&) = delete;
CppFunction(CppFunction&&) noexcept = default;
CppFunction& operator=(CppFunction&&) = default;
static CppFunction makeFromBoxedKernel(c10::BoxedKernel kernel) {
return CppFunction(
c10::KernelFunction::makeFromBoxedKernel(std::move(kernel)),
/* cpp_signature */ std::nullopt, // not known for boxed functions
/* schema */ nullptr);
}
static CppFunction makeFallthrough() {
return makeFromBoxedKernel(c10::BoxedKernel::makeFallthrough());
}
static CppFunction makeNamedNotSupported() {
return makeFromBoxedKernel(c10::BoxedKernel::makeNamedNotSupported());
}
template <c10::BoxedKernel::BoxedKernelFunction* func>
static CppFunction makeFromBoxedFunction() {
return makeFromBoxedKernel(c10::BoxedKernel::makeFromFunction<func>());
}
// Variant that takes in a boxed kernel function with a plumbed
// DispatchKeySet. See Note [Plumbing Keys Through The Dispatcher] for
// details.
template <c10::BoxedKernel::BoxedKernelFunction_withDispatchKeys* func>
static CppFunction makeFromBoxedFunction() {
return makeFromBoxedKernel(c10::BoxedKernel::makeFromFunction<func>());
}
template <class KernelFunctor>
static CppFunction makeFromBoxedFunctor(
std::unique_ptr<KernelFunctor> kernelFunctor) {
return makeFromBoxedKernel(
c10::BoxedKernel::makeFromFunctor(std::move(kernelFunctor)));
}
template <
typename FuncPtr,
std::enable_if_t<
c10::guts::is_function_type<FuncPtr>::value,
std::nullptr_t> = nullptr>
static CppFunction makeFromUnboxedFunction(FuncPtr* f) {
return CppFunction(f);
}
template <
typename FuncPtr,
std::enable_if_t<
c10::is_compile_time_function_pointer<FuncPtr>::value,
std::nullptr_t> = nullptr>
static CppFunction makeFromUnboxedFunction(FuncPtr f) {
return CppFunction(f);
}
CppFunction&& debug(std::string d) && {
debug_ = std::move(d);
return std::move(*this);
}
private:
std::optional<c10::DispatchKey> dispatch_key_;
c10::KernelFunction func_;
std::optional<c10::impl::CppSignature> cpp_signature_;
std::unique_ptr<c10::FunctionSchema> schema_;
std::string debug_;
// The "setter" for dispatch_key_
template <typename Func>
friend CppFunction dispatch(c10::DispatchKey, Func&&);
// The only class which actually pulls out values from CppFunction (does so
// destructively, felt too lazy to write accessors that I don't even
// want users to use)
friend class Library;
CppFunction(
c10::KernelFunction func,
std::optional<c10::impl::CppSignature> cpp_signature,
std::unique_ptr<c10::FunctionSchema> schema);
};
template <typename Func>
inline CppFunction dispatch(c10::DispatchKey k, Func&& raw_f) {
CppFunction f(std::forward<Func>(raw_f));
if (k == c10::DispatchKey::CatchAll) {
f.dispatch_key_ = std::nullopt;
} else {
f.dispatch_key_ = k;
}
return f;
}
template <typename Func>
inline CppFunction dispatch(c10::DeviceType type, Func&& raw_f) {
auto deviceTypeToDispatchKey = [](c10::DeviceType t) {
switch (t) {
// This list is synchronized with the k-constants in c10/core/DeviceType.h
case c10::DeviceType::CPU:
return c10::DispatchKey::CPU;
case c10::DeviceType::CUDA:
return c10::DispatchKey::CUDA;
case c10::DeviceType::IPU:
return c10::DispatchKey::IPU;
case c10::DeviceType::XLA:
return c10::DispatchKey::XLA;
case c10::DeviceType::Lazy:
return c10::DispatchKey::Lazy;
case c10::DeviceType::XPU:
return c10::DispatchKey::XPU;
case c10::DeviceType::MPS:
return c10::DispatchKey::MPS;
case c10::DeviceType::Meta:
return c10::DispatchKey::Meta;
case c10::DeviceType::HIP:
return c10::DispatchKey::HIP;
case c10::DeviceType::MAIA:
return c10::DispatchKey::MAIA;
case c10::DeviceType::HPU:
return c10::DispatchKey::HPU;
case c10::DeviceType::MTIA:
return c10::DispatchKey::MTIA;
case c10::DeviceType::PrivateUse1:
return c10::DispatchKey::PrivateUse1;
default:
TORCH_CHECK(
false,
"Device type ",
t,
" cannot be overloaded at dispatch time, "
"please file a bug report explaining what you were trying to do.");
}
};
return dispatch(deviceTypeToDispatchKey(type), std::forward<Func>(raw_f));
}
inline c10::FunctionSchema schema(const char* str, c10::AliasAnalysisKind k, bool allow_typevars=false) {
c10::FunctionSchema s = torch::jit::parseSchema(str, /*allow_typevars*/allow_typevars);
s.setAliasAnalysis(k);
return s;
}
inline c10::FunctionSchema schema(const char* s, bool allow_typevars=false) {
return schema(s, c10::AliasAnalysisKind::FROM_SCHEMA, allow_typevars);
}
inline c10::FunctionSchema&& schema(c10::FunctionSchema&& s) {
return std::move(s);
}
namespace detail {
inline std::variant<c10::OperatorName, c10::FunctionSchema> constructSchemaOrName(
c10::FunctionSchema&& s) {
return std::move(s);
}
inline std::variant<c10::OperatorName, c10::FunctionSchema> constructSchemaOrName(
c10::OperatorName&& n) {
return std::move(n);
}
inline std::variant<c10::OperatorName, c10::FunctionSchema>
constructSchemaOrName(const char* str) {
auto s = torch::jit::parseSchemaOrName(str);
if (std::holds_alternative<c10::FunctionSchema>(s)) {
std::get<c10::FunctionSchema>(s).setAliasAnalysis(
c10::AliasAnalysisKind::FROM_SCHEMA);
}
return s;
}
class TorchLibraryInit;
} // namespace detail
// Note [Selective build]
// ~~~~~~~~~~~~~~~~~~~~~~
// In some settings, especially mobile, it is important to avoid compiling any
// references to functions that you aren't actually going to use, so that they
// can be eliminated by the linker. We call this capability "selective build".
//
// A very easy way to implement selective build which results in a lot of
// boilerplate is to just add ifdef's around every registration call, but this
// means you have to write a lot of extra lines of code at every registration
// site, and it also means you have to define some munging scheme to map
// operators to macros.
//
// Instead of doing this, we have a different mechanism centered around the
// concept of a SelectiveStr. A selective name is like a const char* string,
// except it also carries at compile time a boolean saying whether or not a
// registration should actually happen or not. We then have extra overloads
// which bypass registration entirely if a selective name is disabled. We do a
// constexpr test to see if a operator should be enabled or not; this is
// currently implemented in ATen/core/op_registration/op_allowlist.h
namespace detail {
// dummy class for non selected custom torchbind classes
class ClassNotSelected {
public:
ClassNotSelected& def_pickle(...) {
return *this;
}
ClassNotSelected& def(...) {
return *this;
}
};
// A SelectiveStr is like a const char*, except that it also comes
// with a type brand that says whether or not the name is enabled or
// not. If the string is disabled, then (at compile time) we DON'T generate
// a registration call for it. This class is not intended to be called
// directly; use TORCH_SELECTIVE_NAME or TORCH_SELECTIVE_SCHEMA macros below
// to create it.
template <bool enabled>
class SelectiveStr {
public:
constexpr explicit SelectiveStr(const char* name) : name_(name) {}
constexpr operator const char*() {
return name_;
}
private:
const char* name_;
};
#define TORCH_SELECTIVE_CLASS(n) \
torch::detail::SelectiveStr<c10::impl::custom_class_allowlist_check(n)>(n)
#define TORCH_SELECTIVE_NAME(n) \
torch::detail::SelectiveStr<c10::impl::op_allowlist_check(n)>(n)
#define TORCH_SELECTIVE_SCHEMA(n) \
torch::detail::SelectiveStr<c10::impl::schema_allowlist_check(n)>(n)
} // namespace detail
class TORCH_API Library final {
public:
enum Kind {
DEF, // from TORCH_LIBRARY (no qualifier)
IMPL,
FRAGMENT,
};
Library(
Kind kind,
std::string ns,
std::optional<c10::DispatchKey> k,
const char* file,
uint32_t line);
Library(const Library&) = delete;
Library& operator=(const Library&) = delete;
Library(Library&&) = default;
Library& operator=(Library&&) = default;
~Library() = default;
// Some notes about the API design here. We had the following constraints:
//
// - We need to support multiple "types" of arguments for schema and
// functions (e.g., unnamed lambda types, regular functions, const char*,
// fully instantiated schemas)
// - We don't want to write exponentially many overloads
// - We don't want to rely on implicit conversion to a common type,
// because the C++ compiler will only be willing to do a single
// implicit conversion (reducing the set of valid types which you
// can invoke with); also error messages are worse when an implicit
// conversion is not selected (as the compiler will not explain
// why it didn't select an implicit conversion; this is different
// from overloads where it will explain each candidate overload and
// why it didn't apply)
//
// To solve all of these constraints at the same time, we use a trick taken
// from the pybind11 library: template over the argument in the user visible
// API, and inside of the templated function explicitly call an overloaded
// function to resolve the argument to a real type. You get the good error
// messages from overloads, but at the same time you only need to write the
// overload for any given argument type once.
template <typename Schema>
Library& def(
Schema&& raw_schema,
const std::vector<at::Tag>& tags = {},
_RegisterOrVerify rv = _RegisterOrVerify::REGISTER) & {
c10::FunctionSchema s = schema(std::forward<Schema>(raw_schema));
return _def(std::move(s), nullptr, tags, rv);
}
Library& set_python_module(const char* pymodule, const char* context = "") {
python_module_ = {pymodule, context};
return *this;
}
Library& impl_abstract_pystub(const char* pymodule, const char* context = "") {
return set_python_module(pymodule, context);
}
template <typename NameOrSchema, typename Func>
Library& def(NameOrSchema&& raw_name_or_schema, Func&& raw_f,
const std::vector<at::Tag>& tags = {}) & {
CppFunction f(std::forward<Func>(raw_f));
return _def(
detail::constructSchemaOrName(
::std::forward<NameOrSchema>(raw_name_or_schema)),
::std::move(f), tags);
}
template <typename Name, typename Func>
Library& impl(
Name name,
Func&& raw_f,
_RegisterOrVerify rv = _RegisterOrVerify::REGISTER) & {
// TODO: need to raise an error when you impl a function that has a
// catch all def
#if defined C10_MOBILE
CppFunction f(std::forward<Func>(raw_f), NoInferSchemaTag());
#else
CppFunction f(std::forward<Func>(raw_f));
#endif
return _impl(name, std::move(f), rv);
}
#if defined C10_MOBILE
// Note: This overload is needed only for C10_MOBILE, since the automatically
// defined copy constructor for the CppFunction doesn't have the additional
// NoInferSchemaTag argument. We define the overload for the impl() function
// to accept a CppFunction&& argument. The already constructed CppFunction
// object may or may not have the inferred schema, but it doesn't matter
// for our purposes since if it already has the inferred schema, then we
// might as well just pass it through directly.
//
template <typename Name>
Library& impl(Name name, CppFunction&& raw_f) & {
// TODO: need to raise an error when you impl a function that has a
// catch all def
CppFunction f(std::forward<CppFunction>(raw_f));
return _impl(name, std::move(f));
}
#endif
// Helper for getting an OperatorName for a const char*. You probably
// don't need this.
c10::OperatorName _resolve(const char* name) const;
template <typename Name, typename Dispatch, typename Func>
Library& impl(Name name, Dispatch&& key, Func&& raw_f) & {
return impl(
name, dispatch(std::forward<Dispatch>(key), std::forward<Func>(raw_f)));
}
template <typename Name, typename Func>
Library& impl_UNBOXED(Name /*name*/, Func* /*raw_f*/) & {
static_assert(
c10::guts::false_t<Func>(),
".impl_UNBOXED(...) was removed. Please use .impl(...) instead.");
return *this;
}
// These overloads cover cases when a SelectiveStr (see Note [Selective
// build]) has been disabled at compile time. In that case, don't generate
// any code referencing the passed in functions at all.
Library& def(detail::SelectiveStr<false>, const std::vector<at::Tag>& tags [[maybe_unused]] = {}) & {
return *this;
}
Library& def(detail::SelectiveStr<true> raw_schema, const std::vector<at::Tag>& tags = {}) & {
return def(raw_schema.operator const char*(), tags);
}
template <typename Func>
Library& def(detail::SelectiveStr<false>, Func&& /*raw_f*/, const std::vector<at::Tag>& tags [[maybe_unused]] = {}) & {
return *this;
}
template <typename Func>
Library& def(detail::SelectiveStr<true> raw_name_or_schema, Func&& raw_f, const std::vector<at::Tag>& tags = {}) & {
return def(
raw_name_or_schema.operator const char*(), std::forward<Func>(raw_f), tags);
}
template <typename Func>
// NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
Library& impl(detail::SelectiveStr<false>, Func&& /*raw_f*/) & {
return *this;
}
template <typename Dispatch, typename Func>
Library& impl(
detail::SelectiveStr<false>,
// NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
Dispatch&& /*key*/,
// NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
Func&& /*raw_f*/) & {
return *this;
}
template <typename Func>
Library& impl_UNBOXED(
detail::SelectiveStr<false> /*name*/,
Func* /*raw_f*/) & {
static_assert(
c10::guts::false_t<Func>(),
".impl_UNBOXED(...) was removed. Please use .impl(...) instead.");
return *this;
}
template <typename Func>
Library& impl(detail::SelectiveStr<true> name, Func&& raw_f) & {
return impl(name.operator const char*(), std::forward<Func>(raw_f));
}
template <typename Dispatch, typename Func>
Library& impl(
detail::SelectiveStr<true> name,
Dispatch&& key,
Func&& raw_f) & {
return impl(
name.operator const char*(),
std::forward<Dispatch>(key),
std::forward<Func>(raw_f));
}
template <typename Func>
Library& impl_UNBOXED(
detail::SelectiveStr<true> /*name*/,
Func* /*raw_f*/) & {
static_assert(
c10::guts::false_t<Func>(),
".impl_UNBOXED(...) was removed. Please use .impl(...) instead.");
return *this;
}
template <typename Func>
Library& fallback(Func&& raw_f) & {
CppFunction f((std::forward<Func>(raw_f)));
return _fallback(std::move(f));
}
template <class CurClass>
inline torch::class_<CurClass> class_(const std::string& className);
// These overloads enable the use of selective build on classes registered
// within a library. The API is the same as before with 1 minor change.
// Instead of m.class_<foo>("foo") you instead do
// m.class_<foo>(TORCH_SELECTIVE_CLASS("foo"))
template <class CurClass>
inline torch::class_<CurClass> class_(detail::SelectiveStr<true> className);
template <class CurClass>
inline detail::ClassNotSelected class_(detail::SelectiveStr<false> className);
// De-registers all registrations created with this Library
void reset();
private:
Kind kind_;
std::optional<std::string> ns_;
std::optional<c10::DispatchKey> dispatch_key_;
std::optional<std::pair<const char*, const char*>> python_module_;
const char* file_;
uint32_t line_;
std::vector<c10::RegistrationHandleRAII> registrars_;
friend class detail::TorchLibraryInit;
// Non-user visible actual implementations of functions. These aren't
// public because we only implement & qualifier and not && qualifier
Library& _def(
c10::FunctionSchema&& schema,
c10::OperatorName* out_name = nullptr,
const std::vector<at::Tag>& tags = {},
_RegisterOrVerify rv = _RegisterOrVerify::REGISTER) &;
Library& _def(
std::variant<c10::OperatorName, c10::FunctionSchema>&&,
CppFunction&& f,
const std::vector<at::Tag>& tags = {}) &;
Library& _impl(
const char* name,
CppFunction&& f,
_RegisterOrVerify rv = _RegisterOrVerify::REGISTER) &;
Library& _fallback(CppFunction&& f) &;
at::OperatorName _parseNameForLib(const char* name_str) const;
};
namespace detail {
class TorchLibraryInit final {
private:
using InitFn = void(Library&);
Library lib_;
public:
TorchLibraryInit(
Library::Kind kind,
InitFn* fn,
const char* ns,
std::optional<c10::DispatchKey> k,
const char* file,
uint32_t line)
: lib_(kind, ns, k, file, line) {
fn(lib_);
}
};
} // namespace detail
} // namespace torch
// NB: The EXACT NAMING of the initializer functions (e.g.,
// TORCH_LIBRARY_init_aten) matters for the code analyzer;
// see the regexes at tools/code_analyzer/run_analyzer.sh
#define TORCH_LIBRARY(ns, m) \
static void TORCH_LIBRARY_init_##ns(torch::Library&); \
static const torch::detail::TorchLibraryInit TORCH_LIBRARY_static_init_##ns( \
torch::Library::DEF, \
&TORCH_LIBRARY_init_##ns, \
#ns, \
std::nullopt, \
__FILE__, \
__LINE__); \
void TORCH_LIBRARY_init_##ns(torch::Library& m)
#define TORCH_LIBRARY_FRAGMENT(ns, m) _TORCH_LIBRARY_FRAGMENT(ns, m, C10_UID)
#define _TORCH_LIBRARY_FRAGMENT(ns, m, uid) \
static void C10_CONCATENATE( \
TORCH_LIBRARY_FRAGMENT_init_##ns##_, uid)(torch::Library&); \
static const torch::detail::TorchLibraryInit C10_CONCATENATE( \
TORCH_LIBRARY_FRAGMENT_static_init_##ns##_, uid)( \
torch::Library::FRAGMENT, \
&C10_CONCATENATE(TORCH_LIBRARY_FRAGMENT_init_##ns##_, uid), \
#ns, \
std::nullopt, \
__FILE__, \
__LINE__); \
void C10_CONCATENATE( \
TORCH_LIBRARY_FRAGMENT_init_##ns##_, uid)(torch::Library & m)
// NB: if the dispatch key is not whitelisted, we simply omit the Library
// call entirely
#define TORCH_LIBRARY_IMPL(ns, k, m) _TORCH_LIBRARY_IMPL(ns, k, m, C10_UID)
#define _TORCH_LIBRARY_IMPL(ns, k, m, uid) \
static void C10_CONCATENATE( \
TORCH_LIBRARY_IMPL_init_##ns##_##k##_, uid)(torch::Library&); \
static const torch::detail::TorchLibraryInit C10_CONCATENATE( \
TORCH_LIBRARY_IMPL_static_init_##ns##_##k##_, uid)( \
torch::Library::IMPL, \
(c10::impl::dispatch_key_allowlist_check(c10::DispatchKey::k) \
? &C10_CONCATENATE(TORCH_LIBRARY_IMPL_init_##ns##_##k##_, uid) \
: [](torch::Library&) -> void {}), \
#ns, \
std::make_optional(c10::DispatchKey::k), \
__FILE__, \
__LINE__); \
void C10_CONCATENATE( \
TORCH_LIBRARY_IMPL_init_##ns##_##k##_, uid)(torch::Library & m)
// These are variants of the macros above which are to be used for testing (they
// don't setup the static initializer, so you can control the visibility of
// the allocated library yourself).
//
// DO NOT use these in production code, they are NOT understood by the
// code analyzer and will be incorrectly analyzed in those situations.
#define MAKE_TORCH_LIBRARY(ns) \
torch::Library(torch::Library::DEF, #ns, std::nullopt, __FILE__, __LINE__)
#define MAKE_TORCH_LIBRARY_IMPL(ns, k) \
torch::Library( \
torch::Library::IMPL, \
#ns, \
std::make_optional(c10::DispatchKey::k), \
__FILE__, \
__LINE__)
// Make the custom class API visible, so it is available from
// torch::Library.
#include <torch/custom_class.h>