Merge pull request #186 from ENSTA-U2IS-AI/dev

🐛 Fix `PackedLayerNorm`

Author: alafage

docs/source/conf.py

@@ -20,7 +20,7 @@
     f"{datetime.now().year!s}, Adrien Lafage and Olivier Laurent"
 )
 author = "Adrien Lafage and Olivier Laurent"
-release = "0.5.2"
+release = "0.5.2.post0"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "torch_uncertainty"
-version = "0.5.2"
+version = "0.5.2.post0"
 authors = [
     { name = "ENSTA U2IS", email = "olivier.laurent@ensta-paris.fr" },
     { name = "Adrien Lafage", email = "adrienlafage@outlook.com" },

tests/layers/test_packed.py

@@ -1,6 +1,7 @@
 import pytest
 import torch
 from einops import repeat
+from torch import nn
 
 from torch_uncertainty.layers.functional.packed import (
     packed_in_projection_packed,
@@ -420,13 +421,15 @@ class TestPackedLayerNorm:
     """Testing the PackedGroupNorm layer class."""
 
     def test_one_estimator_forward(self, batched_qkv: torch.Tensor) -> None:
+        layer_norm = nn.LayerNorm(6)
         packed_layer_norm = PackedLayerNorm(
             embed_dim=6,
             num_estimators=1,
             alpha=1,
         )
-        out = packed_layer_norm(batched_qkv)
-        assert out.shape == torch.Size([2, 3, 6])
+        pe_out = packed_layer_norm(batched_qkv)
+        layer_norm_out = layer_norm(batched_qkv)
+        assert torch.allclose(pe_out, layer_norm_out)
 
 
 class TestPackedMultiheadAttention:

torch_uncertainty/layers/packed.py

@@ -704,6 +704,9 @@ def __init__(
         Shape:
             - Input: :math:`(N, *)` where :math:`*` means any number of additional dimensions.
             - Output: :math:`(N, *)` (same shape as input)
+
+        Warnings:
+            This layer is only suitable to replace ``nn.LayerNorm`` when only the last dimension is normalized.
         """
         super().__init__(
             num_groups=num_estimators,
@@ -715,15 +718,18 @@ def __init__(
         )
 
     def forward(self, inputs: Tensor) -> Tensor:
-        x = rearrange(inputs, "b ... h -> b h ...")
+        shapes = {f"d{i}": size for i, size in enumerate(inputs.shape[1:-1])}
+        shape_str = " ".join(shapes.keys())
+
+        x = rearrange(inputs, "b ... h -> (b ...) h")
         x = F.group_norm(
             x,
             self.num_groups,
             self.weight,
             self.bias,
             self.eps,
         )
-        return rearrange(x, "b h ... -> b ... h")
+        return rearrange(x, f"(b {shape_str}) h -> b {shape_str} h", **shapes)
 
 
 class PackedMultiheadAttention(nn.Module):