Torchscript and C++ interface

[ipcamit]

From the tests/export_test.py file, I could see that for C++ interface I need to use the AOTInductor interface. Problem is that last time I checked, it could only generate forward passes, and not the gradients. Is it possible to use the TorcScript exported model in C++ (using libtorch interface)? I guess the graph is captured, but for C++ runtime, I will need to link against some OpenEquivariance shared library (like in case of torch scatter) ?

TorchScript can provide access to full forwrad-backward graph in C++.

[vbharadwaj-bk]

Yes, you should be able to use TorchScript for this (we have a test for that as well in tests/export_test.py), although we are still working on adding the test for the C++ interface.

And that is also correct; like torch_scatter, you will have to link against an .so file, since PyTorch does not package extension code. We should probably expose this with a better name, but if you run

import openequivariance
print(openequivariance.extlib.torch_module.__file__)

You will get the path to the .so file you need to link.

[ipcamit]

Thanks for the quick reply, I will give it a go.

[vbharadwaj-bk]

@ipcamit We made some updates to iron out some bugs and added a test that shows you how to load a scripted model from start to finish, assuming you have CMake on your system; would you like to try reinstalling the library and running the test? We can get this closed if it works.

pytest tests/export_test.py -k test_jitscript_cpp_interface

File:https://github.com/PASSIONLab/OpenEquivariance/blob/main/tests/export_test.py

[ipcamit]

Sure, I will try in this Thursday, and let you know.

…

On Mon, 2 Jun, 2025, 22:01 Vivek Bharadwaj, ***@***.***> wrote: *vbharadwaj-bk* left a comment (PASSIONLab/OpenEquivariance#107) <#107 (comment)> @ipcamit <https://github.com/ipcamit> We made some updates to iron out some bugs and added a test that shows you how to build an extension from start to finish, assuming you have CMake on your system; would you like to try reinstalling the library and running the test? We can get this closed if it works. pytest tests/export_test.py -k test_jitscript_cpp_interface File: https://github.com/PASSIONLab/OpenEquivariance/blob/main/tests/export_test.py — Reply to this email directly, view it on GitHub <#107 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABTZ2D65HMTASBBZGXDJR6L3BUFXRAVCNFSM6AAAAAB6ISYFLCVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDSMZTGIYDONBQGY> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[ipcamit]

Hi, I think we can close this issue for now. I am little occupied at the moment. Perhaps if it does not work I can ask to open this issue again. Thank you for the prompt replies.

[vbharadwaj-bk]

Sure thing - let us know, and we have documented a new function that gets you the .so path: https://passionlab.github.io/OpenEquivariance/api/#openequivariance.torch_ext_so_path