Layer

Input dimension

Output dimension

Input

-

[84, 84, 5]

-

Lifting

nn.Linear

[84, 84, 5]

[84, 84, 32]

Reshape

Permute

[84, 84, 32]

[32, 84, 84]

Up-padding

F.pad

[32, 84, 84]

[32, 89, 89]

Fourier

Fourier

[32, 89, 89]

[32, 89, 89]

Fourier

Fourier

[32, 89, 89]

[32, 89, 89]

Fourier

Fourier

[32, 89, 89]

[32, 89, 89]

R-Fourier

R-Fourier

[32, 89, 89]

[32, 89, 89]

De-padding

De-padding

[32, 89, 89]

[32, 84, 84]

Projection

MLP

[32, 84, 84]

[2, 84, 84]

Output

MLP

[2, 84, 84]

[2, 84, 84]