Address coderabbitai review comments for AUCMLoss and tests

Signed-off-by: Shubham Chandravanshi <shubham.chandravanshi378@gmail.com>

Author: shubham-61969

monai/losses/aucm_loss.py

@@ -113,6 +113,9 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         input = input.flatten()
         target = target.flatten()
 
+        if input.numel() == 0:
+            raise ValueError("Input and target must contain at least one element.")
+
         if not torch.all((target == 0) | (target == 1)):
             raise ValueError("Target must contain only binary values (0 or 1)")
 
@@ -175,7 +178,10 @@ def _global_mean(self, tensor: torch.Tensor, mask: torch.Tensor) -> torch.Tensor
         Returns:
             Scalar tensor representing the global mean.
         """
-        return (tensor * mask).mean()
+        masked = tensor * mask
+        if masked.numel() == 0:
+            return torch.zeros((), dtype=tensor.dtype, device=tensor.device)
+        return masked.mean()
 
     def _class_mean(self, tensor: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
         """

tests/losses/test_aucm_loss.py

@@ -13,37 +13,63 @@
 
 import unittest
 
+import numpy as np
 import torch
 from parameterized import parameterized
 
 from monai.losses import AUCMLoss
 from tests.test_utils import test_script_save
 
-FIXED_INPUT = torch.tensor([[1.0], [2.0]])
-FIXED_TARGET = torch.tensor([[1.0], [0.0]])
-
-EXPECTED_V1 = 1.25
-EXPECTED_V2 = 5.0
+TEST_CASES = [
+    # small deterministic cases (with expected values)
+    ("v1", torch.tensor([[1.0], [2.0]]), torch.tensor([[1.0], [0.0]]), 1.25),
+    ("v2", torch.tensor([[1.0], [2.0]]), torch.tensor([[1.0], [0.0]]), 5.0),
+]
 
 
 class TestAUCMLoss(unittest.TestCase):
-    """Test cases for AUCMLoss."""
+    """Unit tests for AUCMLoss covering correctness, edge cases, and scriptability."""
 
     @parameterized.expand([("v1",), ("v2",)])
     def test_versions(self, version):
         """Test AUCMLoss with different versions."""
         loss_fn = AUCMLoss(version=version)
-        input = torch.randn(32, 1, requires_grad=True)
+        pred = torch.randn(32, 1, requires_grad=True)
         target = torch.randint(0, 2, (32, 1)).float()
-        loss = loss_fn(input, target)
+        loss = loss_fn(pred, target)
         self.assertIsInstance(loss, torch.Tensor)
         self.assertEqual(loss.ndim, 0)
 
-    @parameterized.expand([("v1", EXPECTED_V1), ("v2", EXPECTED_V2)])
-    def test_known_values(self, version, expected):
+    @parameterized.expand(TEST_CASES)
+    def test_known_values(self, version, pred, target, expected):
         """Test AUCMLoss against fixed manually computed values."""
-        loss = AUCMLoss(version=version)(FIXED_INPUT, FIXED_TARGET)
-        self.assertAlmostEqual(loss.item(), expected, places=5)
+        loss = AUCMLoss(version=version)(pred, target)
+        np.testing.assert_allclose(loss.detach().cpu().numpy(), expected, atol=1e-5, rtol=1e-5)
+
+    @parameterized.expand([("v1",), ("v2",)])
+    def test_high_dimensional(self, version):
+        """Test AUCMLoss with higher dimensional preds (e.g., segmentation)."""
+        loss_fn = AUCMLoss(version=version)
+
+        pred = torch.randn(2, 1, 8, 8, requires_grad=True)
+        target = torch.randint(0, 2, (2, 1, 8, 8)).float()
+
+        loss = loss_fn(pred, target)
+
+        self.assertIsInstance(loss, torch.Tensor)
+        self.assertEqual(loss.ndim, 0)
+
+    def test_imbalanced(self):
+        """Test AUCMLoss with highly imbalanced targets."""
+        loss_fn = AUCMLoss(version="v1")
+
+        pred = torch.randn(32, 1)
+        target = torch.zeros(32, 1)
+        target[0] = 1.0  # only one positive
+
+        loss = loss_fn(pred, target)
+
+        self.assertIsInstance(loss, torch.Tensor)
 
     def test_invalid_version(self):
         """Test that invalid version raises ValueError."""
@@ -57,54 +83,54 @@ def test_invalid_imratio(self):
         with self.assertRaises(ValueError):
             AUCMLoss(imratio=-0.1)
 
-    def test_invalid_input_shape(self):
-        """Test that invalid input shape raises ValueError."""
+    def test_invalid_pred_shape(self):
+        """Test that invalid pred shape raises ValueError."""
         loss_fn = AUCMLoss()
-        input = torch.randn(32, 2)  # Wrong channel
+        pred = torch.randn(32, 2)  # Wrong channel
         target = torch.randint(0, 2, (32, 1)).float()
         with self.assertRaises(ValueError):
-            loss_fn(input, target)
+            loss_fn(pred, target)
 
     def test_invalid_target_shape(self):
         """Test that invalid target shape raises ValueError."""
         loss_fn = AUCMLoss()
-        input = torch.randn(32, 1)
+        pred = torch.randn(32, 1)
         target = torch.randint(0, 2, (32, 2)).float()  # Wrong channel
         with self.assertRaises(ValueError):
-            loss_fn(input, target)
+            loss_fn(pred, target)
 
     def test_insufficient_dimensions(self):
         """Test that tensors with insufficient dimensions raise ValueError."""
         loss_fn = AUCMLoss()
-        input = torch.randn(32)  # 1D tensor
+        pred = torch.randn(32)  # 1D tensor
         target = torch.randint(0, 2, (32, 1)).float()
         with self.assertRaises(ValueError):
-            loss_fn(input, target)
+            loss_fn(pred, target)
 
     def test_shape_mismatch(self):
         """Test that mismatched shapes raise ValueError."""
         loss_fn = AUCMLoss()
-        input = torch.randn(32, 1)
+        pred = torch.randn(32, 1)
         target = torch.randint(0, 2, (16, 1)).float()
         with self.assertRaises(ValueError):
-            loss_fn(input, target)
+            loss_fn(pred, target)
 
     def test_non_binary_target(self):
         """Test that non-binary target values raise ValueError."""
         loss_fn = AUCMLoss()
-        input = torch.randn(32, 1)
+        pred = torch.randn(32, 1)
         target = torch.tensor([[0.5], [1.0], [2.0], [0.0]] * 8)  # 32x1, still non-binary
         with self.assertRaises(ValueError):
-            loss_fn(input, target)
+            loss_fn(pred, target)
 
     def test_backward(self):
         """Test that gradients can be computed."""
         loss_fn = AUCMLoss()
-        input = torch.randn(32, 1, requires_grad=True)
+        pred = torch.randn(32, 1, requires_grad=True)
         target = torch.randint(0, 2, (32, 1)).float()
-        loss = loss_fn(input, target)
+        loss = loss_fn(pred, target)
         loss.backward()
-        self.assertIsNotNone(input.grad)
+        self.assertIsNotNone(pred.grad)
 
     def test_script_save(self):
         """Test that the loss can be saved as TorchScript."""