From 243b6f4efe51b97114185977044b62887512ed51 Mon Sep 17 00:00:00 2001
From: lucidrains <lucidrains@gmail.com>
Date: Sun, 29 Oct 2023 10:46:11 -0700
Subject: [PATCH] last tweak; give a little time overlap when constituting time
 tokens

---
 iTransformer/iTransformer2D.py | 8 ++++++--
 setup.py                       | 2 +-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/iTransformer/iTransformer2D.py b/iTransformer/iTransformer2D.py
index 3508fc8..b88946f 100644
--- a/iTransformer/iTransformer2D.py
+++ b/iTransformer/iTransformer2D.py
@@ -202,9 +202,13 @@ def __init__(
             nn.LayerNorm(dim)
         )
 
+        time_kernel_size = lookback_len // num_time_tokens
+
         self.to_time_tokens = nn.Sequential(
-            Rearrange('b v (t n) -> b v t n', t = num_time_tokens),
-            nn.Linear(lookback_len // num_time_tokens, dim),
+            Rearrange('b v n -> (b v) 1 n'),
+            nn.ConstantPad1d((time_kernel_size, 0), value = 0.),
+            nn.Conv1d(1, dim, time_kernel_size * 2),
+            Rearrange('(b v) d t -> b v t d', v = num_variates),
             nn.LayerNorm(dim)
         )
 
diff --git a/setup.py b/setup.py
index 89028d9..f6cfc74 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'iTransformer',
   packages = find_packages(exclude=[]),
-  version = '0.3.2',
+  version = '0.3.3',
   license='MIT',
   description = 'iTransformer - Inverted Transformer Are Effective for Time Series Forecasting',
   author = 'Phil Wang',