Call ATen kernel from Python?

Hi,
I’ve compiled an ATan kernel using the instructions from the README. I’m wondering if it’s possible to call that kernel from Python, or if that requires writing a separate Python C extension.

Thanks!

This discussion might be relevant:

1 Like

Thank you, works great!

I guess this method is good for custom functions. For the functions already included in PyTorch, a python binding should be generated at compilation. Although this is not the case for some of them at present time I believe…

Are you doing your own custom function with cuda?
Did you manage to compile it with jit?

Ya, using a custom function. https://github.com/malmaud/pytorch_cuda_extension demonstrates the successful attempt.

1 Like