Fix 3D decomposed relative position embedding for anisotropic inputs

monai/networks/blocks/attention_utils.py

@@ -114,13 +114,13 @@ def add_decomposed_rel_pos(
         r_q = q.reshape(batch, q_h, q_w, q_d, dim)
         rel_h = torch.einsum("bhwdc,hkc->bhwdk", r_q, rh)
         rel_w = torch.einsum("bhwdc,wkc->bhwdk", r_q, rw)
-        rel_d = torch.einsum("bhwdc,wkc->bhwdk", r_q, rd)
+        rel_d = torch.einsum("bhwdc,dkc->bhwdk", r_q, rd)
 
         attn = (
             attn.view(batch, q_h, q_w, q_d, k_h, k_w, k_d)
-            + rel_h[:, :, :, :, None, None]
-            + rel_w[:, :, :, None, :, None]
-            + rel_d[:, :, :, None, None, :]
+            + rel_h[:, :, :, :, :, None, None]
+            + rel_w[:, :, :, :, None, :, None]
+            + rel_d[:, :, :, :, None, None, :]
         ).view(batch, q_h * q_w * q_d, k_h * k_w * k_d)
 
     return attn

tests/networks/blocks/test_selfattention.py

@@ -67,6 +67,24 @@ def test_ill_arg(self):
         with self.assertRaises(ValueError):
             SABlock(hidden_size=620, num_heads=8, dropout_rate=0.4)
 
+    # anisotropic 3D shapes with all-distinct, permuted dims: a cubic input_size hides
+    # the decomposed rel-pos axis handling because every misplaced axis still broadcasts.
+    @parameterized.expand([[(4, 8, 16)], [(16, 8, 4)], [(2, 3, 4)], [(4, 3, 2)], [(2, 4, 3)]])
+    @skipUnless(has_einops, "Requires einops")
+    def test_decomposed_rel_pos_anisotropic_3d(self, input_size):
+        hidden_size = 120
+        net = SABlock(
+            hidden_size=hidden_size,
+            num_heads=6,
+            dropout_rate=0.1,
+            rel_pos_embedding=RelPosEmbedding.DECOMPOSED,
+            input_size=input_size,
+        )
+        seq_len = input_size[0] * input_size[1] * input_size[2]
+        with eval_mode(net):
+            result = net(torch.randn(2, seq_len, hidden_size))
+            self.assertEqual(result.shape, (2, seq_len, hidden_size))
+
     def test_rel_pos_embedding_with_flash_attention(self):
         with self.assertRaises(ValueError):
             SABlock(