Added dice loss

docs/src/python/nn/losses.rst

@@ -16,4 +16,5 @@ Loss Functions
    mse_loss
    nll_loss
    smooth_l1_loss
-   triplet_loss
+   triplet_loss
+   dice_loss

python/mlx/nn/losses.py

@@ -274,6 +274,43 @@ def triplet_loss(
     return _reduce(loss, reduction)
 
 
+def dice_loss(
+    predictions: mx.array,
+    targets: mx.array,
+    epsilon: float = 1e-6,
+    reduction: str = "mean",
+) -> mx.array:
+    """
+    Computes the multiclass Dice loss.
+
+    Args:
+        predictions (array): The predicted probabilities. Shape: [batch, classes, ...]
+        targets (array): The target values in one-hot encoded form. Shape: [batch, classes, ...]
+        epsilon (float, optional): Small constant to avoid division by zero. Default: 1e-6.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          'none' | 'mean' | 'sum'. Default: 'mean'.
+
+    Returns:
+        array: The computed multiclass Dice loss.
+    """
+    # Ensure the shapes are correct
+    if predictions.shape != targets.shape:
+        raise ValueError(
+            f"Shape of predictions {predictions.shape} does not match targets shape {targets.shape}."
+        )
+
+    intersection = mx.sum(
+        predictions * targets, axis=[0, *range(2, len(targets.shape))]
+    )
+    cardinality = mx.sum(predictions + targets, axis=[0, *range(2, len(targets.shape))])
+    dice_score = 2.0 * intersection / (cardinality + epsilon)
+
+    dice_loss = 1 - dice_score
+
+    # Apply reduction
+    return _reduce(dice_loss, reduction)
+
+
 def _reduce(loss: mx.array, reduction: str = "none"):
     if reduction == "mean":
         return mx.mean(loss)

python/tests/test_nn.py

@@ -265,6 +265,37 @@ def test_triplet_loss(self):
         expected_sum = mx.sum(expected_none)
         self.assertTrue(mx.allclose(losses_sum, expected_sum))
 
+    def test_dice_loss(self):
+        # Create a simple test case
+        predictions = mx.array([[0.7, 0.2, 0.1], [0.1, 0.6, 0.3]])
+        targets = mx.array([[1, 0, 0], [0, 1, 0]])
+
+        # Expected Dice loss
+        expected_none = mx.array([0.22222265, 0.3333337, 1.0])
+        expected_mean = mx.mean(expected_none)
+        expected_sum = mx.sum(expected_none)
+
+        # Test with reduction 'none'
+        losses_none = nn.losses.dice_loss(predictions, targets, reduction="none")
+        self.assertTrue(
+            mx.allclose(losses_none, expected_none, atol=1e-5),
+            "Test case failed for dice_loss --reduction='none'",
+        )
+
+        # Test with reduction 'mean'
+        losses_mean = nn.losses.dice_loss(predictions, targets, reduction="mean")
+        self.assertTrue(
+            mx.allclose(losses_mean, expected_mean, atol=1e-5),
+            "Test case failed for dice_loss --reduction='mean'",
+        )
+
+        # Test with reduction 'sum'
+        losses_sum = nn.losses.dice_loss(predictions, targets, reduction="sum")
+        self.assertTrue(
+            mx.allclose(losses_sum, expected_sum, atol=1e-5),
+            "Test case failed for dice_loss --reduction='sum'",
+        )
+
     def test_gelu(self):
         inputs = [1.15286231, -0.81037411, 0.35816911, 0.77484438, 0.66276414]