Removed all stale backwards kernel code

Also match the gradient output to the input, in terms of memory layout

tests/test_attention.py

@@ -289,7 +289,7 @@ class TestNeighborhoodAttentionS2(unittest.TestCase):
         v_gpu.requires_grad = True
 
         out_gpu = att_gpu(q_gpu, k_gpu, v_gpu)
-        out_grad = torch.randn(out_gpu.shape, dtype=torch.float32, device="cuda:0").to(memory_format=torch.channels_last)
+        out_grad = torch.randn(out_gpu.shape, dtype=torch.float32, device="cuda:0")
         time_backward_start = torch.cuda.Event(enable_timing=True)
         time_backward_end = torch.cuda.Event(enable_timing=True)
 

torch_harmonics/csrc/attention/attention.cuh

@@ -49,30 +49,3 @@ std::tuple<at::Tensor,at::Tensor,at::Tensor> s2_attention_bwd_dkvq_cuda(at::Tens
                                          at::Tensor psi_col_idx,
                                          at::Tensor psi_row_off,
                                          int nlon_in, int nlat_out, int nlon_out);
-
-torch::Tensor s2_attention_bwd_dq_cuda(at::Tensor kx,
-                                       at::Tensor vx,
-                                       at::Tensor qy,
-                                       at::Tensor dy,
-                                       at::Tensor quad_weights,
-                                       at::Tensor psi_col_idx,
-                                       at::Tensor psi_row_off,
-                                       int nlon_in, int nlat_out, int nlon_out);
-
-torch::Tensor s2_attention_bwd_dk_cuda(at::Tensor kx,
-                                       at::Tensor vx,
-                                       at::Tensor qy,
-                                       at::Tensor dy,
-                                       at::Tensor quad_weights,
-                                       at::Tensor psi_col_idx,
-                                       at::Tensor psi_row_off,
-                                       int nlon_in, int nlat_out, int nlon_out);
-
-torch::Tensor s2_attention_bwd_dv_cuda(at::Tensor kx,
-                                       at::Tensor vx,
-                                       at::Tensor qy,
-                                       at::Tensor dy,
-                                       at::Tensor quad_weights,
-                                       at::Tensor psi_col_idx,
-                                       at::Tensor psi_row_off,
-                                       int nlon_in, int nlat_out, int nlon_out);

torch_harmonics/csrc/attention/attention_interface.cu

@@ -33,10 +33,6 @@
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("forward", &s2_attention_fwd_cuda, "(Local) Attention on S2");
-  m.def("backward_dk", &s2_attention_bwd_dk_cuda, "(Local) Attention gradient on S2 (gradient for k)");
-  m.def("backward_dv", &s2_attention_bwd_dv_cuda, "(Local) Attention gradient on S2 (gradient for v)");
-  m.def("backward_dq", &s2_attention_bwd_dq_cuda,
-        "(Local) Attention gradient on S2 (gradient for q)");
   m.def("backward_dkvq", &s2_attention_bwd_dkvq_cuda, "(Local) Attention gradient on S2 (gradient for k,v,&q)");
 }