Adjust execution order of activation and masking in MaskedDiceLoss

[ytl0623]

Fixes #8655

Description

1. Modified MaskedDiceLoss.forward to ensure masked regions result in 0.0 probability instead of 0.5.
2. Updated unit tests with new expected values, as the fix mathematically changes the loss result by removing the erroneous background contribution.

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

MaskedDiceLoss.forward now applies sigmoid, softmax, and other activations to inputs before masking, then temporarily disables these activations while calling the underlying spatial-weighted path, and finally restores the original activation states. This ensures activations are applied to logits before masking occurs, rather than after. Corresponding test expected values were updated to reflect the corrected computation order.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title accurately describes the main change: reordering activation and masking in MaskedDiceLoss to fix the masked region probability issue.
Description check	✅ Passed	Description covers the core changes (activation before masking), references the linked issue (#8655), explains test updates, and marks the change as non-breaking.
Linked Issues check	✅ Passed	The PR implements the solution to #8655 by applying sigmoid before masking, ensuring masked regions yield 0.0 probability instead of 0.5.
Out of Scope Changes check	✅ Passed	Changes are scoped to the MaskedDiceLoss activation/masking order fix and corresponding test updates; no extraneous modifications detected.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

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

✨ Finishing touches

• 📝 Generate docstrings

---

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

🤖 Fix all issues with AI agents

In `@monai/losses/dice.py`:
- Around line 256-283: The forward block mutates instance flags (self.sigmoid,
self.softmax, self.other_act) before calling spatial_weighted, which is not
thread-safe; instead avoid changing instance state by computing local variables
for the activated input and any activation flags (e.g., use
local_sigmoid/local_softmax/local_other_act or simply pass the already-activated
input) and call spatial_weighted with those locals (or with the activated input
and any needed override parameters) rather than flipping
self.sigmoid/self.softmax/self.other_act; remove the try/finally restore logic
and ensure spatial_weighted signature or call site (spatial_weighted) accepts
the local overrides or only the prepared input to eliminate instance mutation
and race conditions.

🧹 Nitpick comments (3)

monai/losses/dice.py (2)

263-263: Add stacklevel=2 to warning.

Per static analysis (B028): Set stacklevel=2 so the warning points to the caller rather than this line.

Fix

-                warnings.warn("single channel prediction, `softmax=True` ignored.")
+                warnings.warn("single channel prediction, `softmax=True` ignored.", stacklevel=2)

---

257-268: Code duplication with DiceLoss.forward activation handling.

Lines 257-268 mirror DiceLoss.forward (lines 154-165). Consider extracting a _apply_activations helper to reduce duplication.

tests/losses/test_masked_dice_loss.py (1)

115-115: Class name mismatch.

Class TestDiceLoss tests MaskedDiceLoss. Consider renaming to TestMaskedDiceLoss.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

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

📥 Commits

Reviewing files that changed from the base of the PR and between 57fdd59 and 092990d.

📒 Files selected for processing (2)

• monai/losses/dice.py
• tests/losses/test_masked_dice_loss.py

🧰 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:

• tests/losses/test_masked_dice_loss.py
• monai/losses/dice.py

🪛 Ruff (0.14.11)

monai/losses/dice.py

263-263: No explicit stacklevel keyword argument found

Set stacklevel=2

(B028)

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

🔇 Additional comments (1)

tests/losses/test_masked_dice_loss.py (1)

30-30: Updated expected values reflect activation-before-masking fix.

The test values (0.333333, 0.301128, 0.579184) correctly correspond to the implementation change where sigmoid is applied before masking. The code confirms that activation layers execute at lines 257–268, then masked loss computation follows at line 279.

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}