removing a bug with none normalization mode and reverting to L1 normalization

Author: bonevbs

torch_harmonics/attention/attention.py

@@ -47,6 +47,7 @@
 class AttentionS2(nn.Module):
     """
     (Global) attention on the 2-sphere.
+
     Parameters
     -----------
     in_channels: int
@@ -67,6 +68,13 @@ class AttentionS2(nn.Module):
         number of dimensions for interior inner product in the attention matrix (corresponds to kdim in MHA in PyTorch)
     out_channels: int, optional
         number of dimensions for interior inner product in the attention matrix (corresponds to vdim in MHA in PyTorch)
+
+    Reference
+    ---------
+    Bonev, B., Rietmann, M., Paris, A., Carpentieri, A., & Kurth, T. (2025).
+    "Attention on the Sphere."
+    Advances in Neural Information Processing Systems (NeurIPS).
+    https://arxiv.org/abs/2505.11157
     """
 
     def __init__(
@@ -210,6 +218,13 @@ class NeighborhoodAttentionS2(nn.Module):
         number of dimensions for interior inner product in the attention matrix (corresponds to vdim in MHA in PyTorch)
     optimized_kernel: Optional[bool]
         Whether to use the optimized kernel (if available)
+
+    Reference
+    ---------
+    Bonev, B., Rietmann, M., Paris, A., Carpentieri, A., & Kurth, T. (2025).
+    "Attention on the Sphere."
+    Advances in Neural Information Processing Systems (NeurIPS).
+    https://arxiv.org/abs/2505.11157
     """
 
     def __init__(

torch_harmonics/disco/convolution.py

@@ -92,10 +92,6 @@ def _normalize_convolution_tensor_s2(
         If basis_norm_mode is not one of the supported modes.
     """
 
-    # exit here if no normalization is needed
-    if basis_norm_mode == "none":
-        return psi_vals
-
     # reshape the indices implicitly to be ikernel, out_shape[0], in_shape[0], in_shape[1]
     idx = torch.stack([psi_idx[0], psi_idx[1], psi_idx[2] // in_shape[1], psi_idx[2] % in_shape[1]], dim=0)
 
@@ -131,9 +127,10 @@ def _normalize_convolution_tensor_s2(
             # compute the 1-norm
             # scale[ik, ilat] = torch.sqrt(torch.sum(psi_vals[iidx].abs().pow(2) * q[iidx]))
             if basis_norm_mode == "modal":
-                if ik != 0:
-                    bias[ik, ilat] = torch.sum(psi_vals[iidx] * q[iidx])
-                scale[ik, ilat] = torch.sqrt(torch.sum((psi_vals[iidx] - bias[ik, ilat]).abs().pow(2) * q[iidx]))
+                # if ik != 0:
+                #     bias[ik, ilat] = torch.sum(psi_vals[iidx] * q[iidx])
+                # scale[ik, ilat] = torch.sqrt(torch.sum((psi_vals[iidx] - bias[ik, ilat]).abs().pow(2) * q[iidx]))
+                scale[ik, ilat] = torch.sum((psi_vals[iidx] - bias[ik, ilat]).abs() * q[iidx])
             else:
                 scale[ik, ilat] = torch.sum((psi_vals[iidx] - bias[ik, ilat]).abs() * q[iidx])