accept image fmaps

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'soft-moe-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.1',
+  version = '0.0.2',
   license='MIT',
   description = 'Soft MoE - Pytorch',
   author = 'Phil Wang',

soft_moe_pytorch/soft_moe.py

@@ -3,7 +3,7 @@
 import torch.nn.functional as F
 from torch import nn, einsum, Tensor
 
-from einops import rearrange
+from einops import rearrange, pack, unpack
 
 # helper functions
 
@@ -75,6 +75,12 @@ def forward(self, x):
         d - feature dimension
         """
 
+        is_image = x.ndim == 4
+
+        if is_image:
+            x = rearrange(x, 'b d h w -> b h w d')
+            x, ps = pack([x], 'b * d')
+
         # following Algorithm 1, with the normalization they proposed, but with scaling of both (the now popular rmsnorm + gamma)
 
         x = self.norm(x)
@@ -104,4 +110,8 @@ def forward(self, x):
         out = rearrange(out, 'e b s d -> b (e s) d')
         out = einsum('b s d, b n s -> b n d', out, combine_weights)
 
+        if is_image:
+            out, = unpack(out, ps, 'b * d')
+            out = rearrange(out, 'b h w d -> b d h w')
+
         return out