fix transformerblock

[yang-ze-kang]

Fixes # monai/networks/blocks/transformerblock.py

Description

When "with_cross_attention==False", there is no need to initialize "CrossAttentionBlock" in "init"; otherwise, it will introduce unnecessary parameters to the model and may potentially cause some errors.

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

Instantiation of cross-attention components in monai/networks/blocks/transformerblock.py is now conditional: norm_cross_attn and cross_attn are created only when with_cross_attention is True. When False, those attributes are not instantiated and a pre-load hook is registered to drop any state_dict entries related to cross_attn and norm_cross_attn during loading. The forward path remains gated by with_cross_attention. No exported or public entity declarations were altered.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Pre-merge checks and finishing touches

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.
Title Check	❓ Inconclusive	The title is too generic and does not clearly convey the main change of conditional cross-attention initialization in TransformerBlock, so it fails to inform readers what is being fixed.	Please update the title to something more descriptive, for example “Conditionally initialize cross-attention in TransformerBlock when with_cross_attention=False,” so it clearly reflects the key change.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Description Check	✅ Passed	The description follows the repository template by including a summary of the fix, a clear description of the change, and a types of changes checklist, making it largely complete and aligned with expectations.

✨ Finishing touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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]

Actionable comments posted: 1

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

📥 Commits

Reviewing files that changed from the base of the PR and between 53382d8 and 4cd46de.

📒 Files selected for processing (1)

• monai/networks/blocks/transformerblock.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (1)

**/*.py

⚙️ CodeRabbit configuration file

Review the Python code for quality and correctness. Ensure variable names adhere to PEP8 style guides, are sensible and informative in regards to their function, though permitting simple names for loop and comprehension variables. Ensure routine names are meaningful in regards to their function and use verbs, adjectives, and nouns in a semantically appropriate way. Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings. Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these. Suggest any enhancements for code improving efficiency, maintainability, comprehensibility, and correctness. Ensure new or modified definitions will be covered by existing or new unit tests.

Files:

• monai/networks/blocks/transformerblock.py

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (13)

• GitHub Check: flake8-py3 (pytype)
• GitHub Check: flake8-py3 (mypy)
• GitHub Check: quick-py3 (macOS-latest)
• GitHub Check: min-dep-pytorch (2.6.0)
• GitHub Check: flake8-py3 (codeformat)
• GitHub Check: quick-py3 (ubuntu-latest)
• GitHub Check: quick-py3 (windows-latest)
• GitHub Check: packaging
• GitHub Check: min-dep-pytorch (2.7.1)
• GitHub Check: min-dep-pytorch (2.8.0)
• GitHub Check: min-dep-pytorch (2.5.1)
• GitHub Check: min-dep-os (macOS-latest)
• GitHub Check: min-dep-os (windows-latest)

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

monai/networks/blocks/transformerblock.py (1)

94-98: Use the public pre-hook registration.

Switch to register_load_state_dict_pre_hook to avoid leaning on the private helper.

-            self._register_load_state_dict_pre_hook(_drop_cross_attn_keys)
+            self.register_load_state_dict_pre_hook(_drop_cross_attn_keys)

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

📥 Commits

Reviewing files that changed from the base of the PR and between 4cd46de and 2b1a37c.

📒 Files selected for processing (1)

• monai/networks/blocks/transformerblock.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (1)

**/*.py

⚙️ CodeRabbit configuration file

Review the Python code for quality and correctness. Ensure variable names adhere to PEP8 style guides, are sensible and informative in regards to their function, though permitting simple names for loop and comprehension variables. Ensure routine names are meaningful in regards to their function and use verbs, adjectives, and nouns in a semantically appropriate way. Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings. Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these. Suggest any enhancements for code improving efficiency, maintainability, comprehensibility, and correctness. Ensure new or modified definitions will be covered by existing or new unit tests.

Files:

• monai/networks/blocks/transformerblock.py

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (19)

• GitHub Check: packaging
• GitHub Check: build-docs
• GitHub Check: quick-py3 (ubuntu-latest)
• GitHub Check: quick-py3 (windows-latest)
• GitHub Check: flake8-py3 (codeformat)
• GitHub Check: flake8-py3 (pytype)
• GitHub Check: quick-py3 (macOS-latest)
• GitHub Check: flake8-py3 (mypy)
• GitHub Check: min-dep-os (ubuntu-latest)
• GitHub Check: min-dep-pytorch (2.7.1)
• GitHub Check: min-dep-pytorch (2.6.0)
• GitHub Check: min-dep-pytorch (2.8.0)
• GitHub Check: min-dep-pytorch (2.5.1)
• GitHub Check: min-dep-py3 (3.10)
• GitHub Check: min-dep-py3 (3.11)
• GitHub Check: min-dep-os (windows-latest)
• GitHub Check: min-dep-os (macOS-latest)
• GitHub Check: min-dep-py3 (3.9)
• GitHub Check: min-dep-py3 (3.12)