Questions about TorchScript support

[YutackPark]

Hi, I'm exploring the possibility of using cuEquivariance-Torch in a C++ environment, similar to how e3nn models can be exported via TorchScript. I have a few questions:

1. Can cuEquivariance modules be exported using TorchScript?
2. If it is technically feasible but not yet officially supported, are there any plans about it?

I attempted to use both torch.jit.script and torch.jit.trace. While the former raises errors, the latter produces warnings. If it is unexpected, I'll attach a minimal code to reproduce it. Before diving deeper into debugging, I wanted to confirm if there are any related development plans or known limitations.

I'm aware that TorchScript may be deprecated in the future. However, its replacement, torch.export (https://pytorch.org/docs/stable/export.html), is still marked as unstable, and I have no option.

Lastly, I've noticed that using torch.compile with static tensor shapes can nearly double the performance of e3nn. Could cuEquivariance achieve similar speedups with torch.compile, or is this approach less relevant given its use of optimized custom kernels?

Thanks in advance for your guidance and support!

[mitkotak]

Had a similar question regarding torch.compile, does it make sense to put a torch.compile decorator over the FX fallback ? Happy to put in a PR

[borisfom]

@YutackPark : cuEquivariance should be JIT-exportable. We do have jit.trace() unit test in place for underlying cuequivariance_ops_torch, not for this repo yet.
torch.jit.script will fail if you use it on cuEquivariance modules directly - unfortunately, script() restrictions on arguments of exported methods go against API usability and we consider use case of cuEquivariance module being top-level exported module not very common.
If you do wish to script() cuEquivariance module and not the enclosing model, please write a simple wrapper that would not use variable number of arguments or keyword-only arguments with defaults, and script() it. If that fails - please do let us know and submit a repro case.

[borisfom]

@YutackPark : actually, we were able to fix script() compatibility (with changes to the API).
Please check it out, it's merged already: #40

[Yangxinsix]

Hi Boris,

I tried the newest version that fixed torch.compile() and script() compatibility. Now torch.compile works fine but script is still not working. The problem seems come from some API issues in cuequivariance_ops_torch, which is not included in this repo.

Is there any feasible way to circumvent using variable number of arguments or keyword-only arguments in cuequivariance_ops_torch or if there is a newer version?

Thank you for your support!

[borisfom]

@Yangxinsix : thanks for the input! Please watch out for cuequivariance_ops_torch updated package, should come out soon.