Pytorch has (at least) two different jit compile Interfaces: torch.jit.script and the torch.compile.
Both have different performance characteristics and limitations. Example: Torchscript will bail on unsupported operations (so you can change them), torch.compile would only compile the parts it can and use python for the rest.
We should for computer intensive operations mostly elementwise (signal models such as epg, functionals, ...) support at least one of them.
I would say a good goal would be to support torch.compile on all forwards (and adjoints) and test for it. It would be best if we try to limit graph breaks (not sure yet how we could test that though)
I am not sure if torch.compile will have to retrace the FourierOperator.forward if we create a new Operator instance with the same settings (needs testing!)
Additionally, if some functions benefit from torchscript (over torch.compile), we might want to always use it for these and add the decorators.
@ckolbPTB : Small heads up: if I remember correctly, you moved the small functions (RF, gradient, ...) in the epg code to small classes. When I wrote the original code, this causes recompilation if using torchscript. Each class instance was compiled separately. Using functions, pytorch compiled the function only once.
Using torch.jit.script was a massive speed up for this code for longer sequences. This also is the reasons for some of the a little bit strange code constructs as I avoided all unsupported operations (at that time).
Pytorch has (at least) two different jit compile Interfaces: torch.jit.script and the torch.compile.
Both have different performance characteristics and limitations. Example: Torchscript will bail on unsupported operations (so you can change them), torch.compile would only compile the parts it can and use python for the rest.
We should for computer intensive operations mostly elementwise (signal models such as epg, functionals, ...) support at least one of them.
I would say a good goal would be to support torch.compile on all forwards (and adjoints) and test for it. It would be best if we try to limit graph breaks (not sure yet how we could test that though)
I am not sure if torch.compile will have to retrace the FourierOperator.forward if we create a new Operator instance with the same settings (needs testing!)
Additionally, if some functions benefit from torchscript (over torch.compile), we might want to always use it for these and add the decorators.
@ckolbPTB : Small heads up: if I remember correctly, you moved the small functions (RF, gradient, ...) in the epg code to small classes. When I wrote the original code, this causes recompilation if using torchscript. Each class instance was compiled separately. Using functions, pytorch compiled the function only once. Using torch.jit.script was a massive speed up for this code for longer sequences. This also is the reasons for some of the a little bit strange code constructs as I avoided all unsupported operations (at that time).