Register a python object with TORCH_LIBRARY_IMPL

C++
1

Hello.

Before I ask, I apologize for my lack of knowledge as this is my first time asking questions about open source projects.

I’ve seen in the manual how to use the TORCH_LIBRARY_IMPL macro to register a kernel pointer that maps to a specific backend key.
However, this is only available for the C++ implementation.

Is it possible to register a python object inside the TORCH_LIBRARY_IMPL C++ macro?

Inevitably, I need to call the python object implemented at the python level like a kernel, which I’ve found can only be registered by using library.impl or equivalent decorator at the python level.

So, while analyzing the source code, I found at::impl::PythonOpRegistrationTrampoline, which seemed to can possible to register a Python module in TORCH_LIBRARY_IMPL in C++.

Would it be okay to use this? and if it works, would it be okay to ask review the code for upstream?

2

Are you looking for something like torch.library — PyTorch 2.5 documentation

3

torch.library you mentioned allows you to register a python object as a kernel for a specific backend at the python level.
But, I’m looking for a way to register Python objects at C++ level. so, I’m looking for a way to register a Python kernel using a C++ macro like TORCH_LIBRARY_IMPL.

4

It certainly seems possible as you point out, but out of curiosity, why do you need to support this?

For a single op definition, it is okay to have a partially initialized fragment whether that is in cpp or in python. So someone should be able to import your cpp extension and later import the python module that registers the python kernel impl.