Fix layernorm.

Author: ameya98

src/e3tools/nn/_layer_norm.py

@@ -116,7 +116,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
             if is_scalar:
                 # For scalar irreps (l=0, p=1), use standard layer norm
-                field_norm = F.layer_norm(field_view, [dim], None, None, self.eps)
+                field_norm = F.layer_norm(field_view, [mul, dim], None, None, self.eps)
                 # Flatten back for concatenation
                 field_out = field_norm.reshape(*batch_dims, size)
             else:

tests/test_layer_norm.py

@@ -24,6 +24,7 @@ def test_equivariance(irreps_in: str):
     [
         "0e + 1o",
         "32x0e + 1o + 2e",
+        "0e + 4x1o + 5e",
         "3x1o + 2x2o",
         "8x1o + 2x2o + 1x3o",
         "3x1o + 2x2o + 1x3o + 1x4o",
@@ -43,7 +44,7 @@ def test_layer_norm_compiled(irreps_in: str, seed: int, batch_size: int = 8):
     assert torch.allclose(output, output_compiled)
 
 
-@pytest.mark.parametrize("irreps_in", ["0e + 1o", "32x0e + 1o + 2e", "3x1o + 2x2o"])
+@pytest.mark.parametrize("irreps_in", ["0e + 1o", "32x0e + 1o + 2e", "0e + 4x1o + 5e", "3x1o + 2x2o"])
 def test_layer_norm(irreps_in: str):
     irreps_in = e3nn.o3.Irreps(irreps_in)
     layer = LayerNorm(irreps_in)
@@ -53,9 +54,9 @@ def test_layer_norm(irreps_in: str):
     input = irreps_in.randn(-1)
     output = layer(input)
 
-    for _, ir, field in unpack_irreps(output, irreps_in):
+    for mul, ir, field in unpack_irreps(output, irreps_in):
         sq_norms = field.norm(dim=-1, keepdim=True).pow(2).sum(dim=-1).mean(dim=-1)
-        if ir.l == 0 and ir.p == 1:
+        if ir.l == 0 and ir.p == 1 and mul == 1:
             assert torch.allclose(sq_norms, torch.as_tensor([0.0]))
         else:
             assert torch.allclose(sq_norms, torch.as_tensor([1.0]))