1616@withCUDA
1717@pytest .mark .parametrize ("cross_channel" , [False , True ])
1818@pytest .mark .parametrize ("alpha" , [0.0 , 0.5 , 1.0 ])
19+ @pytest .mark .parametrize ("context_size" , [1 , 37 , 1024 ])
1920def test_forward (
2021 monkeypatch : pytest .MonkeyPatch ,
2122 cross_channel : bool ,
2223 alpha : float ,
24+ context_size : int ,
2325 device : torch .device ,
2426) -> None :
2527 monkeypatch .setattr (
@@ -38,20 +40,20 @@ def test_forward(
3840
3941 # Past and past-and-future covariates:
4042 x_context = TableTensor .from_tensor (
41- torch .randn (1024 , 2 , device = device ),
43+ torch .randn (context_size , 2 , device = device ),
4244 columns = ["x1" , "x2" ],
4345 )
4446 x_query = TableTensor .from_tensor (torch .randn (120 , 0 , device = device ))
4547
4648 # Target variates:
4749 y_context = TableTensor .from_tensor (
48- torch .randn (1024 , 3 , device = device ),
50+ torch .randn (context_size , 3 , device = device ),
4951 columns = ["y1" , "y2" , "y3" ],
5052 )
5153
5254 with (
5355 nullcontext ()
54- if cross_channel or alpha < 1
56+ if cross_channel or ( alpha < 1 and context_size >= 632 )
5557 else pytest .warns (RuntimeWarning , match = "Historical covariates" )
5658 ):
5759 out = model (x_context , y_context , x_query , alpha = alpha )
@@ -61,3 +63,33 @@ def test_forward(
6163 assert "y1" in out .numerical_columns
6264 assert "y2" in out .numerical_columns
6365 assert "y3" in out .numerical_columns
66+
67+ if context_size < 512 :
68+ padded_x_context = TableTensor .from_tensor (
69+ torch .nn .functional .pad (
70+ x_context .numerical ,
71+ pad = (0 , 0 , 512 - context_size , 0 ),
72+ value = float ("nan" ),
73+ ),
74+ columns = x_context .numerical_columns ,
75+ )
76+ padded_y_context = TableTensor .from_tensor (
77+ torch .nn .functional .pad (
78+ y_context .numerical ,
79+ pad = (0 , 0 , 512 - context_size , 0 ),
80+ value = float ("nan" ),
81+ ),
82+ columns = y_context .numerical_columns ,
83+ )
84+ with (
85+ nullcontext ()
86+ if cross_channel
87+ else pytest .warns (RuntimeWarning , match = "Historical covariates" )
88+ ):
89+ expected = model (
90+ padded_x_context ,
91+ padded_y_context ,
92+ x_query ,
93+ alpha = alpha ,
94+ )
95+ torch .testing .assert_close (out .numerical , expected .numerical )
0 commit comments