Bbonev/discrete continuous convolutions (#23)

* Adding prototype implementation of disco convolutions

* adding s2convolutions.py

* Somewhat functional first implementation in Triton

* Somewhat functional first implementation in Triton

* Made batched execution work on GPU and refactored precomputation of kernel values and supports.

* Adding the Lobatto grid to the grid selection method in quadrature

* fixing bug in triton kernel when iterating over non-zeros

* Updating with working custom triton implementation of DISCO convolution

* Merged both forward and backward DISCO contraction into a single kernel

* Some cleanup and minor bugfix

* bugfixes to the reference implementation

* adding weights to s2conv

* removing unnecessary imports

* suggestion for torch harmonics, math should be checked thoroughly

* Intermediate working reference implementation for the transpose DISCO convolution. Fixed normalization of kernels

* adjusting cutoff frequency in disco convolution

* fixing transpose DISCO contraction

* moving triton to install_requires

* enabling triton by default

---------

Co-authored-by: Thorsten Kurth <tkurth@nvidia.com>

Author: bonevbs

torch_harmonics/s2_convolutions.py

@@ -194,7 +194,7 @@ def __init__(
         else:
             self.bias = None
 
-    def forward(self, x: torch.Tensor, use_triton_kernel: bool = False) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, use_triton_kernel: bool = True) -> torch.Tensor:
         # pre-multiply x with the quadrature weights
         x = self.quad_weights * x