Register GlobalMutualInformationLoss bin_centers as buffer

[ugbotueferhire]

Fixes #8866.

Description

This PR fixes a device placement bug in GlobalMutualInformationLoss.__init__ where bin_centers was assigned as a plain Python attribute instead of being registered as a buffer.

• Registered bin_centers as a non-persistent buffer in image_dissimilarity.py, ensuring that GlobalMutualInformationLoss now follows normal nn.Module buffer semantics for .to(device) / dtype moves and avoids silent gradient tracking.
• Added a regression test in test_global_mutual_information_loss.py to verify that bin_centers is properly exposed through named_buffers(), does not require gradients, and successfully changes dtype when the module is moved.

Verification: Passed python -m pytest tests/losses/image_dissimilarity/test_global_mutual_information_loss.py -q -k gaussian_bin_centers_registered_buffer and -k ill_opts.

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 registers bin_centers as a non-persistent PyTorch buffer instead of a plain attribute. In __init__, the buffer is initialized to None and populated for Gaussian kernels with expanded dimensions. parzen_windowing_gaussian validates the buffer is defined before use. Two tests confirm Gaussian kernels register a non-trainable buffer that follows dtype and device changes, while b-spline kernels keep bin_centers as None.

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 16.67% 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 change: registering bin_centers as a PyTorch buffer in GlobalMutualInformationLoss.
Description check	✅ Passed	Description is complete, references issue #8866, explains the fix and test coverage, and follows the template structure with appropriate checkboxes marked.
Linked Issues check	✅ Passed	Code changes fully satisfy the objectives from issue #8866: bin_centers is registered as a non-persistent buffer [#8866], validation added in parzen_windowing_gaussian [#8866], and regression tests verify buffer semantics and device placement [#8866].
Out of Scope Changes check	✅ Passed	All changes directly address issue #8866: buffer registration, validation logic, and regression tests. No unrelated modifications detected.

_{✏️ 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]

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/losses/image_dissimilarity.py (1)

202-227: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win

Document the bin_centers buffer.

The __init__ docstring should mention bin_centers as registered module state. Add a note explaining it holds bin centers for Gaussian kernel Parzen windowing (shape: (1, 1, num_bins) for broadcasting). As per coding guidelines, docstrings should describe all variables and state.

🤖 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 202 - 227, Update the
__init__ docstring to document the registered buffer bin_centers: state that
bin_centers is stored as a module buffer (registered via self.register_buffer)
containing the Gaussian-Parzen window bin centers used for intensity Parzen
windowing, its purpose (used by the Gaussian kernel to compute soft
histogram/probability per bin), and its expected shape for broadcasting (1, 1,
num_bins). Mention it is created in the constructor alongside other args
(kernel_type, num_bins, sigma_ratio, etc.) so readers know it is part of the
module state.

🧹 Nitpick comments (1)

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

119-127: ⚡ Quick win

Add docstring and consider shape verification.

The test method lacks a docstring explaining what aspects of buffer registration it verifies. Also consider checking bin_centers.shape equals (1, 1, 16) to confirm the shape expansion. As per coding guidelines, docstrings should be present for all definitions.

📋 Suggested improvements

     def test_gaussian_bin_centers_registered_buffer(self):
+        """Verify bin_centers is a registered buffer with correct dtype/device behavior."""
         loss = GlobalMutualInformationLoss(kernel_type="gaussian", num_bins=16)
 
         self.assertIn("bin_centers", dict(loss.named_buffers()))
         self.assertFalse(loss.bin_centers.requires_grad)
+        self.assertEqual(loss.bin_centers.shape, (1, 1, 16))
 
         loss = loss.to(dtype=torch.float64)

🤖 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 119 - 127, Add a docstring to the test method
test_gaussian_bin_centers_registered_buffer describing that it verifies buffer
registration, dtype preservation on .to(), and expected shape; then add an
assertion that loss.bin_centers.shape == (1, 1, 16) to verify the expanded
shape, and keep the existing checks for presence in named_buffers(),
requires_grad False, and dtype after loss.to(dtype=torch.float64).

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

Inline comments:
In `@monai/losses/image_dissimilarity.py`:
- Around line 236-239: The type annotation for self.bin_centers is declared
unconditionally but the buffer is only created when self.kernel_type ==
"gaussian", so move the annotation inside that conditional: remove the top-level
"self.bin_centers: torch.Tensor" and add the annotation immediately before or
alongside the register_buffer call within the if block that sets self.preterm
and calls self.register_buffer("bin_centers", ...); this ensures
self.bin_centers is only typed/defined when the gaussian branch executes (refer
to the conditional using self.kernel_type, the attribute self.preterm, and the
register_buffer call for locating the code).

In `@tests/losses/image_dissimilarity/test_global_mutual_information_loss.py`:
- Around line 119-127: Update the test_gaussian_bin_centers_registered_buffer to
also assert that the registered buffer moves devices when the module is moved:
create a CUDA device if torch.cuda.is_available(), call loss = loss.to(device),
then assert loss.bin_centers.device == device; keep existing dtype/assertions
and use the same GlobalMutualInformationLoss and bin_centers symbols so the
check validates device placement after .to(device).

---

Outside diff comments:
In `@monai/losses/image_dissimilarity.py`:
- Around line 202-227: Update the __init__ docstring to document the registered
buffer bin_centers: state that bin_centers is stored as a module buffer
(registered via self.register_buffer) containing the Gaussian-Parzen window bin
centers used for intensity Parzen windowing, its purpose (used by the Gaussian
kernel to compute soft histogram/probability per bin), and its expected shape
for broadcasting (1, 1, num_bins). Mention it is created in the constructor
alongside other args (kernel_type, num_bins, sigma_ratio, etc.) so readers know
it is part of the module state.

---

Nitpick comments:
In `@tests/losses/image_dissimilarity/test_global_mutual_information_loss.py`:
- Around line 119-127: Add a docstring to the test method
test_gaussian_bin_centers_registered_buffer describing that it verifies buffer
registration, dtype preservation on .to(), and expected shape; then add an
assertion that loss.bin_centers.shape == (1, 1, 16) to verify the expanded
shape, and keep the existing checks for presence in named_buffers(),
requires_grad False, and dtype after loss.to(dtype=torch.float64).

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 4e9d79f5-6f44-4caa-b664-3f5dcd5cf326

📥 Commits

Reviewing files that changed from the base of the PR and between 0a8d945 and 57b049e.

📒 Files selected for processing (2)

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

[ericspod]

Hi @ugbotueferhire thanks for the contribution. Please look at the coderabbit comments and the DCO fail. The test fails are something I'm working on in another PR.

[ericspod]

Hi @ugbotueferhire this looks fine but please look at the Coderabbit comments and the DCO issue. I've updated so the tests should all work now.

[coderabbitai]

Actionable comments posted: 4

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

Inline comments:
In `@monai/losses/image_dissimilarity.py`:
- Around line 236-241: Remove the leftover merge-conflict markers and make a
single consistent declaration for bin_centers: replace the conflict block with a
single definition that both types the attribute as optional (self.bin_centers:
torch.Tensor | None) and registers it as a buffer
(self.register_buffer("bin_centers", None, persistent=False)); ensure there are
no "<<<<<", ">>>>>" or "=======" markers remaining and that only the
self.bin_centers annotation and the register_buffer call remain.

In `@tests/inferers/test_diffusion_inferer.py`:
- Around line 134-147: Add Google-style docstrings to the new test helper
classes and method: _DiffusersLikeOutput, _DiffusersStyleScheduler, and the step
method inside test_diffusers_style_ddpm_sampler; describe the purpose of each
class, document __init__ parameters prev_sample and pred_original_sample in the
Args section, document what step() accepts (model_output, timestep, sample,
generator, return_dict) and what it returns (either an instance of
_DiffusersLikeOutput or a tuple (prev_sample, pred_original_sample)), and list
any exceptions it may raise in the Raises section; ensure docstrings follow the
Google style (short summary, Args, Returns, Raises) and are added directly above
each class and method definition.
- Around line 132-165: Remove the unresolved merge conflict markers and restore
the intended test function: delete the lines starting with <<<<<<<, ======= and
>>>>>>> and ensure the test_diffusers_style_ddpm_sampler function (including
classes _DiffusersLikeOutput and _DiffusersStyleScheduler, model/scheduler
setup, and assertions) is present once in the file, placed before the
parameterized.expand(TEST_CASES) / `@skipUnless` block so the file parses
correctly and the test runs.

In `@tests/losses/image_dissimilarity/test_global_mutual_information_loss.py`:
- Around line 123-143: Remove the unresolved git conflict markers in the test
file and restore the intended assertions around GlobalMutualInformationLoss and
its bin_centers; specifically, delete the "<<<<<<<", "=======", and ">>>>>>>"
markers and merge the stashed assertions so the test checks that
loss.bin_centers is not None and requires_grad is False, that dtype conversion
via loss.to(dtype=torch.float64) updates loss.bin_centers.dtype, that moving to
CUDA updates loss.bin_centers.device when CUDA is available, and that for
kernel_type="b-spline" loss.bin_centers is None (references:
GlobalMutualInformationLoss, loss.bin_centers,
test_b_spline_bin_centers_exists_as_none).

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 3c6b37a2-8f43-41a6-ae19-67d488a9c6bd

📥 Commits

Reviewing files that changed from the base of the PR and between 57b049e and 4f17f96.

📒 Files selected for processing (3)

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

[coderabbitai]

🧹 Nitpick comments (1)

monai/losses/image_dissimilarity.py (1)

319-320: ⚡ Quick win

Document the new ValueError in the method docstring.

parzen_windowing_gaussian now raises ValueError, but the docstring doesn’t list it in a Raises section.

Proposed docstring update

     def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Parzen windowing with gaussian kernel (adapted from DeepReg implementation)
         Note: the input is expected to range between 0 and 1
         Args:
             img: the shape should be B[NDHW].
+        Raises:
+            ValueError: When ``self.bin_centers`` is not initialized.
         """

As per coding guidelines, “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 319 - 320, The method
parzen_windowing_gaussian now raises ValueError when self.bin_centers is None;
update the function docstring to include a Raises section documenting this
ValueError (e.g., "Raises: ValueError: if bin_centers is None, required for
gaussian parzen windowing."), keep wording consistent with Google-style
docstrings and place it alongside existing Args and Returns entries for
parzen_windowing_gaussian.

🤖 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 319-320: The method parzen_windowing_gaussian now raises
ValueError when self.bin_centers is None; update the function docstring to
include a Raises section documenting this ValueError (e.g., "Raises: ValueError:
if bin_centers is None, required for gaussian parzen windowing."), keep wording
consistent with Google-style docstrings and place it alongside existing Args and
Returns entries for parzen_windowing_gaussian.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 0cca837a-9593-4d80-9e9b-7944831c4706

📥 Commits

Reviewing files that changed from the base of the PR and between 4f17f96 and d3e1571.

📒 Files selected for processing (2)

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

🚧 Files skipped from review as they are similar to previous changes (1)

• tests/losses/image_dissimilarity/test_global_mutual_information_loss.py

[ugbotueferhire]

Hi @ericspod, thanks for the guidance. I addressed the CodeRabbit feedback as suggested

I cleaned up the accidental conflict markers/unrelated file from the branch, rebased onto latest dev, squashed to one signed-off commit, and force-pushed. The focused ruff and pytest checks pass locally.