Prevent half-precision mutual information overflow

[yinkev]

Closes #7622.

Summary

• compute Gaussian Parzen-window weights in float32 for float16 and bfloat16 inputs
• accumulate joint and marginal histograms in float32 with autocast disabled
• preserve loss, weight, probability, device, and gradient dtype contracts after the stable internal computation
• add regressions for many-bin Gaussian weights, large explicit reduced-precision volumes, and float16 AMP/autocast

Root cause

Two reduced-precision operations can become non-finite:

1. with many Gaussian bins, the kernel preterm exceeds the float16 range and can produce NaN weights;
2. the joint-histogram bmm accumulates over all voxels before division, so AMP can overflow the float16 accumulator and return -inf loss with non-finite gradients.

The fix keeps public outputs in their original dtype while performing the numerically sensitive kernel and histogram accumulation in float32. Float64 inputs remain float64.

Validation

• upstream/dev produces non-finite many-bin float16 weights
• under CPU float16 autocast, upstream/dev returns -inf with non-finite gradients on a 48³ volume
• reduced-precision test class: 5 passed across float16, bfloat16, and CPU float16 autocast
• the affected test module passed locally before the review-only test expansion
• Black, isort, Ruff, DCO sign-off, and git diff --check pass

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

GlobalMutualInformationLoss now performs half- and bfloat16 Gaussian Parzen and mutual information calculations using float32 intermediates, disables autocast during accumulation, and restores the original output dtype. New tests verify finite Gaussian outputs, loss values, and gradients for half-precision inputs.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The PR addresses issue #7622 by eliminating non-finite GlobalMutualInformationLoss values during training.
Out of Scope Changes check	✅ Passed	The code changes and tests stay focused on the reported overflow fix and related regression coverage.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Title check	✅ Passed	The title clearly and concisely summarizes the main fix: preventing half-precision mutual information overflow.
Description check	✅ Passed	The description covers the change, root cause, validation, and types of changes, matching the template closely.

✨ 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.}

[coderabbitai]

🧹 Nitpick comments (2)

monai/losses/image_dissimilarity.py (1)

320-344: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Add required Google-style docstrings.

• monai/losses/image_dissimilarity.py#L320-L344: document returned weight/probability tensors and the buffer-related ValueError.

• monai/losses/image_dissimilarity.py#L359-L383: document the returned reduced loss tensor.

• tests/losses/image_dissimilarity/test_global_mutual_information_loss.py#L167-L197: add docstrings for the new test class and methods.

  As per path instructions, “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.”

🤖 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 `@monai/losses/image_dissimilarity.py` around lines 320 - 344, The Gaussian
Parzen windowing implementation in monai/losses/image_dissimilarity.py:320-344
must document the returned weight and probability tensors and the ValueError
raised when bin_centers or preterm is unset; the reduced loss implementation in
monai/losses/image_dissimilarity.py:359-383 must document its returned loss
tensor; and the new test class and methods in
tests/losses/image_dissimilarity/test_global_mutual_information_loss.py:167-197
must receive Google-style docstrings describing their purpose, variables,
returns, and exceptions where applicable.

Source: Path instructions

tests/losses/image_dissimilarity/test_global_mutual_information_loss.py (1)

168-195: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Cover bfloat16 and direct Parzen return contracts.

The implementation supports bfloat16, but these regressions exercise only float16. Run both dtypes and assert weight/probability preserve dtype and device as well as finiteness.

🤖 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/losses/image_dissimilarity/test_global_mutual_information_loss.py`
around lines 168 - 195, Update the tests
test_half_precision_gaussian_weights_with_many_bins_are_finite and
test_half_precision_large_constant_volume_is_finite to run for both float16 and
bfloat16. In the direct parzen_windowing_gaussian assertions, also verify weight
and probability preserve the input dtype and device, while retaining the
existing finiteness checks and loss/gradient contract assertions.

Source: Path instructions

🤖 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.

Nitpick comments:
In `@monai/losses/image_dissimilarity.py`:
- Around line 320-344: The Gaussian Parzen windowing implementation in
monai/losses/image_dissimilarity.py:320-344 must document the returned weight
and probability tensors and the ValueError raised when bin_centers or preterm is
unset; the reduced loss implementation in
monai/losses/image_dissimilarity.py:359-383 must document its returned loss
tensor; and the new test class and methods in
tests/losses/image_dissimilarity/test_global_mutual_information_loss.py:167-197
must receive Google-style docstrings describing their purpose, variables,
returns, and exceptions where applicable.

In `@tests/losses/image_dissimilarity/test_global_mutual_information_loss.py`:
- Around line 168-195: Update the tests
test_half_precision_gaussian_weights_with_many_bins_are_finite and
test_half_precision_large_constant_volume_is_finite to run for both float16 and
bfloat16. In the direct parzen_windowing_gaussian assertions, also verify weight
and probability preserve the input dtype and device, while retaining the
existing finiteness checks and loss/gradient contract assertions.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 0da9611f-61d1-4c73-98fc-b32ddc2ca78b

📥 Commits

Reviewing files that changed from the base of the PR and between 3ee058b and 0ee73aa.

📒 Files selected for processing (2)

• monai/losses/image_dissimilarity.py
• tests/losses/image_dissimilarity/test_global_mutual_information_loss.py

[yinkev]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[yinkev]

Replaced by #9023 to conform to this repository's documented contributor workflow and branch-naming convention. The replacement uses the same signed commit and identical code diff: #9023