Fix bug with empty slice if amount of dimensions matches amount of frequencies, but is not equal to 512

modeling_reve.py

@@ -229,6 +229,19 @@ class FourierEmb4D(nn.Module):
         self.freqs = freqs
         self.increment_time = increment_time
         self.margin = margin
+        harmonics_size = 2 * self.freqs ** 4
+        if self.dimension % 2 != 0:
+            raise ValueError(
+                "Requested embedding size must be even. " \
+                f"Requested: {self.dimension}, Must be even as amount of frequency components in positional encodings " \
+                "is always even due to amount of harmonics being doubled (one for sine and one for cosine component)."
+            )
+        if self.dimension > harmonics_size:
+            raise ValueError(
+                f"Requested embedding size is larger than the maximum possible size for the given frequency range. " \
+                f"Requested: {self.dimension}, Maximum: {harmonics_size}. " \
+                "Please increase the number of frequencies or decrease the embedding size."
+            )
 
     def forward(self, positions_):
         positions = positions_.clone()
@@ -247,9 +260,18 @@ class FourierEmb4D(nn.Module):
         p_w = 2 * math.pi * freqs_w / width
         positions = positions[..., None, None, None, None, :]
         loc = (positions[..., 0] * p_x + positions[..., 1] * p_y + positions[..., 2] * p_z + positions[..., 3] * p_w).view(*U, -1)
-        if self.dimension != 512:  # noqa
-            _, _, hd = loc.shape
-            diff = hd - self.dimension // 2
+        # Positional encoding returns dimension of size self.freqs ** 4 per sine/cosine pair
+        # That makes output dimention 2 * self.freqs ** 4
+        # If requested embedding size of a model is larger final harmonics are cut off
+        _, _, hd = loc.shape
+        diff = hd - self.dimension // 2
+        if diff < 0:
+            raise ValueError(
+                "Requested embedding size is larger than the maximum possible size for the given frequency range. " \
+                f"Requested: {self.dimension}, Maximum: {2 * self.freqs ** 4}. " \
+                "Please increase the number of frequencies or decrease the embedding size."
+            )
+        elif diff > 0:
             loc = loc[:, :, :-diff]
         emb = torch.cat([torch.cos(loc), torch.sin(loc)], dim=-1)
         return emb