handle single conditioning broadcasted across entire sequence of features

Author: lucidrains

mogrifier/mogrifier.py

@@ -4,7 +4,7 @@
 from torch import nn, Tensor
 from torch.nn import Module
 
-from einops import pack, unpack
+from einops import repeat, pack, unpack
 
 # constants
 
@@ -42,7 +42,7 @@ def __init__(
         iters = 5,
         factorize_k: int | None = None,
         dim_hidden: int | None = None,
-        hidden_factorize_k: int | None = None
+        hidden_factorize_k: int | None = None,
     ):
         super().__init__()
         assert iters > 1
@@ -74,6 +74,9 @@ def forward(
     ):
         iters = default(iters, self.iters)
 
+        if inputs.ndim == 3 and hiddens.ndim == 2:
+            hiddens = repeat(hiddens, 'b d -> b n d', n = inputs.shape[-2])
+
         assert inputs.shape[-1] == self.dim
         assert hiddens.shape[-1] == self.dim_hidden
         assert inputs.shape[:-2] == hiddens.shape[:-2]

setup.py

@@ -3,7 +3,7 @@
 setup(
     name = 'mogrifier',
     packages = find_packages(),
-    version = '0.0.4',
+    version = '0.0.5',
     license='MIT',
     description = 'Implementation of Mogrifier circuit from Deepmind',
     long_description_content_type = 'text/markdown',