Fix 3D decomposed relative position embedding for anisotropic inputs

[aymuos15]

The 3D branch of add_decomposed_rel_pos was broken for non-cubic spatial sizes in two ways:

1. rel_d contracted the depth-axis embedding against the width axis ('bhwdc,wkc->bhwdk' instead of 'bhwdc,dkc->bhwdk').
2. The broadcast-add into the (b, qh, qw, qd, kh, kw, kd) attention tensor used 4 leading colons instead of 5 for each rel_* term, so the key axis landed in the wrong slot (e.g. rel_h's kh fell into the kd position).

Both were masked by the only 3D test using a cubic input_size (8, 8, 8): with kh == kw == kd every misplaced axis still broadcasts. For anisotropic volumes the forward pass raised a shape error (or, where conformable, added the wrong positional bias) in SABlock/CrossAttentionBlock with rel_pos_embedding='decomposed'.

Adds a regression test parametrized over several anisotropic 3D input sizes, including permutations of distinct dims so each spatial axis is exercised in a different position.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR fixes the depth-axis relative position embedding computation in 3D attention. The add_decomposed_rel_pos function's einsum contraction for rel_d was corrected from wkc to dkc indexing, with the attention-logit update reformatted across lines. A new parameterized test validates the fix across anisotropic 3D input configurations.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title accurately summarizes the main fix: correcting 3D decomposed relative position embeddings for anisotropic inputs, which is the core change.
Description check	✅ Passed	Description covers all template requirements: explains the two bugs fixed, notes why previous tests missed them, lists new tests added, and correctly marks non-breaking change and new tests.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tests/networks/blocks/test_selfattention.py (1)

251-251: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Pre-existing bug: use_combined_linear is undefined.

This variable is not defined in test_no_extra_weights_if_no_fc - it's only a loop variable in test_script. Test will raise NameError.

Proposed fix

     `@parameterized.expand`([[True], [False]])
-    def test_no_extra_weights_if_no_fc(self, include_fc):
+    def test_no_extra_weights_if_no_fc(self, include_fc, use_combined_linear=True):
         input_param = {

Or update @parameterized.expand to include use_combined_linear values.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/networks/blocks/test_selfattention.py` at line 251, The test fails
because use_combined_linear is referenced but not defined in
test_no_extra_weights_if_no_fc; fix by explicitly defining use_combined_linear
(e.g., insert use_combined_linear = False) inside test_no_extra_weights_if_no_fc
before building the kwargs dictionary that includes "use_combined_linear", or
alternatively update the `@parameterized.expand` on test_no_extra_weights_if_no_fc
to include the use_combined_linear values (matching how test_script parametrizes
it) so the symbol is provided to the test.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@tests/networks/blocks/test_selfattention.py`:
- Line 251: The test fails because use_combined_linear is referenced but not
defined in test_no_extra_weights_if_no_fc; fix by explicitly defining
use_combined_linear (e.g., insert use_combined_linear = False) inside
test_no_extra_weights_if_no_fc before building the kwargs dictionary that
includes "use_combined_linear", or alternatively update the
`@parameterized.expand` on test_no_extra_weights_if_no_fc to include the
use_combined_linear values (matching how test_script parametrizes it) so the
symbol is provided to the test.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: e35d751b-ce63-42f8-b2a3-201db9d1633f

📥 Commits

Reviewing files that changed from the base of the PR and between ddf3795 and 59ce580.

📒 Files selected for processing (2)

• monai/networks/blocks/attention_utils.py
• tests/networks/blocks/test_selfattention.py