From f11047aa54d12f853064e48dbed9fb963a3f10ca Mon Sep 17 00:00:00 2001
From: ayyucedemirbas <a.ayyuced@gmail.com>
Date: Wed, 20 Dec 2023 01:51:44 +0300
Subject: [PATCH 1/2] added dice loss

---
 docs/src/python/nn/losses.rst |  3 ++-
 python/mlx/nn/losses.py       | 37 +++++++++++++++++++++++++++++++++++
 python/tests/test_nn.py       | 37 +++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/docs/src/python/nn/losses.rst b/docs/src/python/nn/losses.rst
index b6a202d4a9..70f5228ec4 100644
--- a/docs/src/python/nn/losses.rst
+++ b/docs/src/python/nn/losses.rst
@@ -16,4 +16,5 @@ Loss Functions
    mse_loss
    nll_loss
    smooth_l1_loss
-   triplet_loss
\ No newline at end of file
+   triplet_loss
+   dice_loss
\ No newline at end of file
diff --git a/python/mlx/nn/losses.py b/python/mlx/nn/losses.py
index 755656e4fa..23fa8408a1 100644
--- a/python/mlx/nn/losses.py
+++ b/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)
diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py
index 17ec5175c0..861b689b66 100644
--- a/python/tests/test_nn.py
+++ b/python/tests/test_nn.py
@@ -265,6 +265,43 @@ 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]
 

From aa9c715c06740372c66fd0438ac5c939de473935 Mon Sep 17 00:00:00 2001
From: ayyucedemirbas <a.ayyuced@gmail.com>
Date: Wed, 20 Dec 2023 02:06:48 +0300
Subject: [PATCH 2/2] added dice_loss

---
 python/tests/test_nn.py | 12 +++---------
 1 file changed, 3 insertions(+), 9 deletions(-)

diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py
index 861b689b66..e61d60e04c 100644
--- a/python/tests/test_nn.py
+++ b/python/tests/test_nn.py
@@ -276,27 +276,21 @@ def test_dice_loss(self):
         expected_sum = mx.sum(expected_none)
 
         # Test with reduction 'none'
-        losses_none = nn.losses.dice_loss(
-            predictions, targets, 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"
-        )
+        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"
-        )
+        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'",