Diffusion Model Encoder has an output layer set in forward method and this leads to problems

[virginiafdez]

Fixes #8577 .

Description

This pull request sets the out layer of DiffusionModelEncoder in the init method. This requires the inclusion of input_shape parameter in the init method to calculate the input dimension to the last linear layer.

The output spatial shape derivation is a bit baroque, but allows for the otherwise not very pleasant odd spatial dimensions at input.

Types of changes

• Non breaking change (fix or new feature that would cause existing functionality to change). The tutorial that was failing will need to provide this input parameter, but input_shape is now defaulted to the notebook dims.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.

[coderabbitai]

Caution

Review failed

The head commit changed during the review from a560756 to 7ca0676.

Walkthrough

• Added input_shape: Sequence[int] = (64, 64) to DiffusionModelEncoder constructor and docstring.
• Imported numpy as np and functools.reduce for dimension computations.
• In init, iteratively halves input_shape per channel level (ceiling) to model downsampling.
• Computes last_dim_flattened = product(halved spatial dims) * channels[-1].
• Initializes a feed-forward output head in init: Linear(last_dim_flattened, 512) → ReLU → Dropout(0.1) → Linear(512, out_channels).
• In forward, assigns a new Sequential to self.out unconditionally, recreating the head each call.
• Public API updated due to new input_shape parameter.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Linked Issues Check	⚠️ Warning	The PR partly addresses issue #8577 by adding input_shape and creating a concrete output head in init, but the forward method still reassigns self.out unconditionally which recreates the head every forward pass; this preserves the original problem of unstable/incorrectly placed parameters and thus does not fully satisfy the linked issue's objectives. Because the head can still be re-created at runtime, the change fails to guarantee the layer's parameters are consistently registered and placed on the correct device.	Remove any reassignment of self.out from the forward method so the head is only constructed in init, compute and store the flattened input dimension in init (validating input_shape), add unit tests that assert self.out parameters appear in model.parameters() and that model.to(device) moves them, and update tutorials/docs to reflect the API change.
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The title directly and accurately identifies the primary issue the PR addresses — that the DiffusionModelEncoder's output layer was being created inside the forward method and causing problems. It is specific and relevant to the changes in this PR and clear enough for a reviewer scanning history.
Out of Scope Changes Check	✅ Passed	All modifications are focused on the output-head initialization and related dimension calculation and the added input_shape constructor parameter; I did not observe unrelated file edits or functionality changes outside the scope of fixing the output-layer initialization. The primary concern remaining is the forward-method reassignment which is scope-related but incorrect.
Description Check	✅ Passed	The PR description follows the repository template: it includes "Fixes #8577", a concise Description that explains the change and rationale, and the Types of changes checklist with test and doc information. There is a minor inconsistency where the Types of changes checkbox marks non-breaking while PR metadata and objectives indicate this introduces a breaking API change; clarify which is correct.

---

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

Caution

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

⚠️ Outside diff range comments (1)

monai/networks/nets/diffusion_model_unet.py (1)

1989-1989: Off-by-one error in final block check

The condition incorrectly uses len(channels) instead of len(channels) - 1, causing incorrect downsampling behavior.

-            is_final_block = i == len(channels)  # - 1
+            is_final_block = i == len(channels) - 1

🧹 Nitpick comments (2)

monai/networks/nets/diffusion_model_unet.py (2)

36-39: Unnecessary imports added

functools.reduce and numpy are only used in DiffusionModelEncoder. Consider moving these imports closer to where they're used or making them conditional imports to avoid unnecessary dependencies for users of DiffusionModelUNet.

---

1887-1887: Missing input_shape in docstring

The new input_shape parameter needs documentation in the docstring.

Add to the Args section:

     Args:
         spatial_dims: number of spatial dimensions.
         in_channels: number of input channels.
         out_channels: number of output channels.
+        input_shape: spatial shape of the input (without batch and channel dims).
         num_res_blocks: number of residual blocks (see _ResnetBlock) per level.

📜 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 6327a86 and b0bb6bf.

📒 Files selected for processing (1)

• monai/networks/nets/diffusion_model_unet.py (5 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/nets/diffusion_model_unet.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: flake8-py3 (pytype)
• GitHub Check: quick-py3 (macOS-latest)
• GitHub Check: flake8-py3 (codeformat)
• GitHub Check: quick-py3 (ubuntu-latest)
• GitHub Check: quick-py3 (windows-latest)
• GitHub Check: flake8-py3 (mypy)
• GitHub Check: build-docs
• GitHub Check: min-dep-pytorch (2.8.0)
• GitHub Check: min-dep-py3 (3.9)
• GitHub Check: min-dep-py3 (3.12)
• GitHub Check: min-dep-os (windows-latest)
• GitHub Check: min-dep-pytorch (2.7.1)
• GitHub Check: min-dep-pytorch (2.5.1)
• GitHub Check: min-dep-pytorch (2.6.0)
• GitHub Check: min-dep-os (ubuntu-latest)
• GitHub Check: min-dep-py3 (3.11)
• GitHub Check: min-dep-os (macOS-latest)
• GitHub Check: min-dep-py3 (3.10)

🔇 Additional comments (3)

monai/networks/nets/diffusion_model_unet.py (3)

2014-2017: Ceil division maintains odd dimensions correctly

The spatial dimension calculation using np.ceil properly handles odd-sized inputs as intended.

---

2019-2021: Output layer correctly moved to init

The MLP head is now properly instantiated during initialization, resolving the device and optimizer issues.

---

1543-1543: Document or remove the new input_shape parameter

Constructor adds input_shape (monai/networks/nets/diffusion_model_unet.py ≈ line 1543) but verification scripts returned no output — manually confirm whether input_shape is referenced anywhere; if unused remove it, otherwise wire it into the class and add/update its parameter docs and tests.

[coderabbitai]

Actionable comments posted: 2

Caution

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

⚠️ Outside diff range comments (1)

monai/networks/nets/diffusion_model_unet.py (1)

1989-1989: Bug: Off-by-one error in final block check

This will cause an index out of bounds error. The condition should be i == len(channels) - 1.

-            is_final_block = i == len(channels)  # - 1
+            is_final_block = i == len(channels) - 1

🧹 Nitpick comments (2)

monai/networks/nets/diffusion_model_unet.py (2)

1907-1907: Document the expected format of input_shape

Add clarification about whether this should include batch/channel dimensions.

-        input_shape: spatial shape of the input (without batch and channel dims).
+        input_shape: spatial shape of the input (H, W) for 2D or (H, W, D) for 3D, excluding batch and channel dims.

---

2019-2021: Type annotation inconsistency

self.out should not be Optional since it's always initialized.

-        self.out: Optional[nn.Module] = nn.Sequential(
+        self.out = nn.Sequential(

📜 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 6327a86 and cbf5dc5.

📒 Files selected for processing (1)

• monai/networks/nets/diffusion_model_unet.py (5 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/nets/diffusion_model_unet.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: build-docs
• GitHub Check: quick-py3 (macOS-latest)
• GitHub Check: flake8-py3 (pytype)
• GitHub Check: quick-py3 (windows-latest)
• GitHub Check: flake8-py3 (mypy)
• GitHub Check: quick-py3 (ubuntu-latest)
• GitHub Check: packaging
• GitHub Check: flake8-py3 (codeformat)
• GitHub Check: min-dep-pytorch (2.6.0)
• GitHub Check: min-dep-pytorch (2.7.1)
• GitHub Check: min-dep-os (windows-latest)
• GitHub Check: min-dep-pytorch (2.8.0)
• GitHub Check: min-dep-os (ubuntu-latest)
• GitHub Check: min-dep-py3 (3.9)
• GitHub Check: min-dep-pytorch (2.5.1)
• GitHub Check: min-dep-os (macOS-latest)
• GitHub Check: min-dep-py3 (3.11)
• GitHub Check: min-dep-py3 (3.12)
• GitHub Check: min-dep-py3 (3.10)

[KumoLiu]

/build