Repository: pyg-team/pytorch_geometric
Pull request #6785 (2023-02-24 09:45:12)

Added the `TemporalEncoding` module

codecov[bot]: # [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6785?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#6785](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6785?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (ec803e1) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/d336d13e09b9740549746800aafa295cc6cfe45d?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (d336d13) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head ec803e1 differs from pull request most recent head a665629. Consider uploading reports for the commit a665629 to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #6785   +/-   ##
=======================================
  Coverage   91.36%   91.37%           
=======================================
  Files         425      425           
  Lines       22936    22950   +14     
=======================================
+ Hits        20956    20970   +14     
  Misses       1980     1980           
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6785?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/\_\_init\_\_.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6785?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/nn/encoding.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6785?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2VuY29kaW5nLnB5) | `100.00% <100.00%> (ø)` | |

: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 6bac441d2..238723791 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
-- Added `pip` caching in CI ([#6462](https://github.com/pyg-team/pytorch_geometric/pull/6462))
+- Added the `TemporalEncoding` module ([#6785](https://github.com/pyg-team/pytorch_geometric/pull/6785))
 - Added CPU-optimized `spmm_reduce` functionality via CSR format ([#6699](https://github.com/pyg-team/pytorch_geometric/pull/6699), [#6759](https://github.com/pyg-team/pytorch_geometric/pull/6759))
 - Added support for the revised version of the `MD17` dataset ([#6734](https://github.com/pyg-team/pytorch_geometric/pull/6734))
 - Added TorchScript support to the `RECT_L` model ([#6727](https://github.com/pyg-team/pytorch_geometric/pull/6727))
diff --git a/torch_geometric/nn/__init__.py b/torch_geometric/nn/__init__.py
index c3cb5d2cf..fb6c56b1c 100644
--- a/torch_geometric/nn/__init__.py
+++ b/torch_geometric/nn/__init__.py
@@ -4,7 +4,7 @@ from .data_parallel import DataParallel
 from .to_hetero_transformer import to_hetero
 from .to_hetero_with_bases_transformer import to_hetero_with_bases
 from .to_fixed_size_transformer import to_fixed_size
-from .encoding import PositionalEncoding
+from .encoding import PositionalEncoding, TemporalEncoding
 from .model_hub import PyGModelHubMixin
 from .summary import summary
 
@@ -27,6 +27,7 @@ __all__ = [
     'to_hetero_with_bases',
     'to_fixed_size',
     'PositionalEncoding',
+    'TemporalEncoding',
     'PyGModelHubMixin',
     'summary',
 ]
diff --git a/torch_geometric/nn/encoding.py b/torch_geometric/nn/encoding.py
index 4bef1838c..cdead3b24 100644
--- a/torch_geometric/nn/encoding.py
+++ b/torch_geometric/nn/encoding.py
@@ -1,9 +1,11 @@
+import math
+
 import torch
 from torch import Tensor
 
 
 class PositionalEncoding(torch.nn.Module):
-    r"""The positional encoding scheme from `"Attention Is All You Need"
+    r"""The positional encoding scheme from the `"Attention Is All You Need"
     <https://arxiv.org/pdf/1706.03762.pdf>`_ paper
 
     .. math::
@@ -54,3 +56,41 @@ class PositionalEncoding(torch.nn.Module):
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}({self.out_channels})'
+
+
+class TemporalEncoding(torch.nn.Module):
+    r"""The time-encoding function from the `"Do We Really Need Complicated
+    Model Architectures for Temporal Networks?"
+    <https://openreview.net/forum?id=ayPPc0SyLv1>`_ paper.
+    :class:`TemporalEncoding` first maps each entry to a vector with
+    monotonically exponentially decreasing values, and then uses the cosine
+    function to project all values to range :math:`[-1, 1]`
+
+    .. math::
+        y_{i} = \cos \left(x \cdot \sqrt{d}^{-(i - 1)/\sqrt{d}} \right)
+
+    where :math:`d` defines the output feature dimension, and
+    :math:`1 \leq i \leq d`.
+
+    Args:
+        out_channels (int): Size :math:`d` of each output sample.
+    """
+    def __init__(self, out_channels: int):
+        super().__init__()
+        self.out_channels = out_channels
+
+        sqrt = math.sqrt(out_channels)
+        weight = 1.0 / sqrt**torch.linspace(0, sqrt, out_channels).view(1, -1)
+        self.register_buffer('weight', weight)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        pass
+
+    def forward(self, x: Tensor) -> Tensor:
+        """"""
+        return torch.cos(x.view(-1, 1) @ self.weight)
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.out_channels})'


Test patch:
diff --git a/test/nn/test_encoding.py b/test/nn/test_encoding.py
index 4106742e3..36278c91d 100644
--- a/test/nn/test_encoding.py
+++ b/test/nn/test_encoding.py
@@ -1,11 +1,22 @@
 import torch
 
-from torch_geometric.nn import PositionalEncoding
+from torch_geometric.nn import PositionalEncoding, TemporalEncoding
+from torch_geometric.testing import withCUDA
 
 
-def test_positional_encoding():
-    encoder = PositionalEncoding(64)
+@withCUDA
+def test_positional_encoding(device):
+    encoder = PositionalEncoding(64).to(device)
     assert str(encoder) == 'PositionalEncoding(64)'
 
-    x = torch.tensor([1.0, 2.0, 3.0])
+    x = torch.tensor([1.0, 2.0, 3.0], device=device)
+    assert encoder(x).size() == (3, 64)
+
+
+@withCUDA
+def test_temporal_encoding(device):
+    encoder = TemporalEncoding(64).to(device)
+    assert str(encoder) == 'TemporalEncoding(64)'
+
+    x = torch.tensor([1.0, 2.0, 3.0], device=device)
     assert encoder(x).size() == (3, 64)