Repository: pyg-team/pytorch_geometric
Pull request #7208 (2023-04-19 10:41:42)

Introduce `type_ptr` argument to `HeteroLayerNorm`

codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7208?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7208](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7208?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (3f893f9) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/9236cb700fa90f8a72511622c2c6df409f3c7b96?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (9236cb7) will **decrease** coverage by `0.47%`.
> The diff coverage is `65.38%`.

```diff
@@            Coverage Diff             @@
##           master    #7208      +/-   ##
==========================================
- Coverage   92.03%   91.56%   -0.47%     
==========================================
  Files         436      436              
  Lines       24050    24067      +17     
==========================================
- Hits        22135    22038      -97     
- Misses       1915     2029     +114     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7208?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/norm/layer\_norm.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7208?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL25vcm0vbGF5ZXJfbm9ybS5weQ==) | `94.44% <60.00%> (-5.56%)` | :arrow_down: |
| [torch\_geometric/profile/benchmark.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7208?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3Byb2ZpbGUvYmVuY2htYXJrLnB5) | `86.79% <68.75%> (-8.21%)` | :arrow_down: |

... and [21 files with indirect coverage changes](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7208/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 91cde2ba5..7bf92c39e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `type_ptr` argument to `HeteroLayerNorm` ([#7208](https://github.com/pyg-team/pytorch_geometric/pull/7208))
 - Added an option to benchmark scripts to write PyTorch profiler results to CSV ([#7114](https://github.com/pyg-team/pytorch_geometric/pull/7114))
 - Added subgraph type sampling option with bidirectional edge support ([#7199](https://github.com/pyg-team/pytorch_geometric/pull/7199), [#7200](https://github.com/pyg-team/pytorch_geometric/pull/7200))
 - Added support for `"any"`-reductions in `scatter` ([#7198](https://github.com/pyg-team/pytorch_geometric/pull/7198))
diff --git a/torch_geometric/nn/norm/layer_norm.py b/torch_geometric/nn/norm/layer_norm.py
index c40d3cc9c..fdaaeeb76 100644
--- a/torch_geometric/nn/norm/layer_norm.py
+++ b/torch_geometric/nn/norm/layer_norm.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import List, Optional, Union
 
 import torch
 import torch.nn.functional as F
@@ -143,6 +143,7 @@ class HeteroLayerNorm(torch.nn.Module):
         mode: str = 'node',
     ):
         super().__init__()
+        assert mode == 'node'
 
         self.in_channels = in_channels
         self.num_types = num_types
@@ -164,16 +165,41 @@ class HeteroLayerNorm(torch.nn.Module):
             torch.nn.init.ones_(self.weight)
             torch.nn.init.zeros_(self.bias)
 
-    def forward(self, x: Tensor, type_vec: Tensor) -> Tensor:
+    def forward(
+        self,
+        x: Tensor,
+        type_vec: OptTensor = None,
+        type_ptr: Optional[Union[Tensor, List[int]]] = None,
+    ) -> Tensor:
         r"""
+        .. note::
+            Either :obj:`type_vec` or :obj:`type_ptr` needs to be specified.
+            In general, relying on :obj:`type_ptr` is more efficient in case
+            the input tensor is sorted by types.
+
         Args:
             x (torch.Tensor): The input features.
-            type_vec (torch.Tensor): A vector that maps each entry to a type.
+            type_vec (torch.Tensor, optional): A vector that maps each entry to
+                a type. (default: :obj:`None`)
+            type_ptr (torch.Tensor or List[int]): A vector denoting the
+                boundaries of types. (default: :obj:`None`)
         """
+        if type_vec is None and type_ptr is None:
+            raise ValueError("Either 'type_vec' or 'type_ptr' must be given")
+
         out = F.layer_norm(x, (self.in_channels, ), None, None, self.eps)
 
         if self.affine:
-            out = out * self.weight[type_vec] + self.bias[type_vec]
+            # TODO Revisit this logic completely as it performs worse than just
+            # operating on a dictionary of tensors
+            # (especially the `type_vec` code path)
+            if type_ptr is not None:
+                h = torch.empty_like(out)
+                for i, (s, e) in enumerate(zip(type_ptr[:-1], type_ptr[1:])):
+                    h[s:e] = out[s:e] * self.weight[i] + self.bias[i]
+                out = h
+            else:
+                out = out * self.weight[type_vec] + self.bias[type_vec]
 
         return out
 


Test patch:
diff --git a/test/nn/norm/test_layer_norm.py b/test/nn/norm/test_layer_norm.py
index 3c17c1670..ba3ac977f 100644
--- a/test/nn/norm/test_layer_norm.py
+++ b/test/nn/norm/test_layer_norm.py
@@ -31,26 +31,39 @@ def test_layer_norm(device, affine, mode):
 @pytest.mark.parametrize('affine', [False, True])
 def test_hetero_layer_norm(device, affine):
     x = torch.randn((100, 16), device=device)
+    expected = LayerNorm(16, affine=affine, mode='node').to(device)(x)
 
     # Test single type:
-    norm = LayerNorm(16, affine=affine, mode='node').to(device)
-    expected = norm(x)
-
     type_vec = torch.zeros(100, dtype=torch.long, device=device)
+    type_ptr = [0, 100]
+
     norm = HeteroLayerNorm(16, num_types=1, affine=affine).to(device)
     assert str(norm) == 'HeteroLayerNorm(16, num_types=1)'
 
     out = norm(x, type_vec)
     assert out.size() == (100, 16)
-    assert torch.allclose(out, expected)
+    assert torch.allclose(out, expected, atol=1e-3)
+    assert torch.allclose(norm(out, type_ptr=type_ptr), expected, atol=1e-3)
+
+    mean = out.mean(dim=-1)
+    std = out.std(unbiased=False, dim=-1)
+    assert torch.allclose(mean, torch.zeros_like(mean), atol=1e-2)
+    assert torch.allclose(std, torch.ones_like(std), atol=1e-2)
 
     # Test multiple types:
-    type_vec = torch.randint(5, (100, ), device=device)
+    type_vec = torch.arange(5, device=device)
+    type_vec = type_vec.view(-1, 1).repeat(1, 20).view(-1)
+    type_ptr = [0, 20, 40, 60, 80, 100]
+
     norm = HeteroLayerNorm(16, num_types=5, affine=affine).to(device)
+    assert str(norm) == 'HeteroLayerNorm(16, num_types=5)'
+
     out = norm(x, type_vec)
     assert out.size() == (100, 16)
+    assert torch.allclose(out, expected, atol=1e-3)
+    assert torch.allclose(norm(out, type_ptr=type_ptr), expected, atol=1e-3)
 
     mean = out.mean(dim=-1)
     std = out.std(unbiased=False, dim=-1)
-    assert torch.allclose(mean, torch.zeros_like(mean), atol=1e-5)
-    assert torch.allclose(std, torch.ones_like(std), atol=1e-5)
+    assert torch.allclose(mean, torch.zeros_like(mean), atol=1e-2)
+    assert torch.allclose(std, torch.ones_like(std), atol=1e-2)