fix: fix confusion matrix float32 problem

Author: edge7-jl

keras/src/metrics/iou_metrics.py

@@ -69,6 +69,7 @@ def __init__(
             name="total_confusion_matrix",
             shape=(num_classes, num_classes),
             initializer=initializers.Zeros(),
+            dtype="float64",
         )
 
     def update_state(self, y_true, y_pred, sample_weight=None):
@@ -131,7 +132,7 @@ def update_state(self, y_true, y_pred, sample_weight=None):
             y_pred,
             self.num_classes,
             weights=sample_weight,
-            dtype="float32",
+            dtype="float64",
         )
 
         return self.total_cm.assign(self.total_cm + current_cm)

keras/src/metrics/iou_metrics_test.py

@@ -1,10 +1,10 @@
 import numpy as np
 import pytest
 
+from keras.src import backend
 from keras.src import layers
 from keras.src import models
 from keras.src import testing
-from keras.src import backend
 from keras.src.metrics import iou_metrics as metrics
 
 
@@ -362,7 +362,8 @@ def confusion_matrix(self, y_true, y_pred, num_classes):
         - num_classes: int, number of classes.
 
         Returns:
-        - conf_matrix: np.ndarray, confusion matrix of shape (num_classes, num_classes).
+        - conf_matrix: np.ndarray, confusion matrix of shape (num_classes,
+                                                              num_classes).
         """
         y_true = np.asarray(y_true)
         y_pred = np.asarray(y_pred)
@@ -374,7 +375,8 @@ def confusion_matrix(self, y_true, y_pred, num_classes):
         conf_matrix = conf_matrix.reshape((num_classes, num_classes))
         return conf_matrix
 
-    def _get_big_chunk(self):
+    @staticmethod
+    def _get_big_chunk():
         np.random.seed(14)
         all_y_true = np.random.choice([0, 1, 2], size=(600, 530, 530))
         # Generate random probabilities for each channel
@@ -389,28 +391,46 @@ def _get_big_chunk(self):
         tmp_pred = np.reshape(all_y_pred_arg, -1)
         return all_y_true, all_y_pred_arg, mean_iou_metric, tmp_true, tmp_pred
 
-    @pytest.mark.skipif(backend.backend() not in ["tensorflow", "torch"], reason="tensorflow/torch only test.")
+    @pytest.mark.skipif(
+        backend.backend() not in ["tensorflow", "torch"],
+        reason="tensorflow/torch only test.",
+    )
     def test_big_chunk_tf_torch(self):
-        all_y_true, all_y_pred_arg, mean_iou_metric_all, tmp_true, tmp_pred = self._get_big_chunk()
+        all_y_true, all_y_pred_arg, mean_iou_metric_all, tmp_true, tmp_pred = (
+            self._get_big_chunk()
+        )
         conf_matrix_from_keras = mean_iou_metric_all.total_cm.numpy()
         # Validate confusion matrices and results
         conf_matrix_manual = self.confusion_matrix(tmp_true, tmp_pred, 3)
-        self.assertTrue(np.array_equal(conf_matrix_from_keras, conf_matrix_manual),
-                        msg='Confusion matrices do not match!')
+        self.assertTrue(
+            np.array_equal(conf_matrix_from_keras, conf_matrix_manual),
+            msg="Confusion matrices do not match!",
+        )
 
     @pytest.mark.skipif(backend.backend() != "jax", reason="jax only test.")
     def test_big_chunk_jax(self):
-        import jax.numpy as jnp
-        all_y_true, all_y_pred_arg, mean_iou_metric_all, tmp_true, tmp_pred = self._get_big_chunk()
-        conf_matrix_from_keras = jnp.array(mean_iou_metric_all.total_cm)
-        # Validate confusion matrices and results
-        conf_matrix_manual = self.confusion_matrix(tmp_true, tmp_pred, 3)
-        self.assertTrue(np.array_equal(conf_matrix_from_keras, conf_matrix_manual),
-                        msg='Confusion matrices do not match!')
-
+        from jax import config
 
+        config.update("jax_enable_x64", False)
 
+        all_y_true, all_y_pred_arg, mean_iou_metric_all, tmp_true, tmp_pred = (
+            self._get_big_chunk()
+        )
+        conf_matrix_from_keras = np.array(mean_iou_metric_all.total_cm)
+        # Validate confusion matrices and results
+        conf_matrix_manual = self.confusion_matrix(tmp_true, tmp_pred, 3)
+        self.assertFalse(
+            np.array_equal(conf_matrix_from_keras, conf_matrix_manual),
+            msg="Confusion matrices do not match!",
+        )
 
+        config.update("jax_enable_x64", True)
+        _, _, mean_iou_metric_all, _, _ = self._get_big_chunk()
+        conf_matrix_from_keras = np.array(mean_iou_metric_all.total_cm)
+        self.assertTrue(
+            np.array_equal(conf_matrix_from_keras, conf_matrix_manual),
+            msg="Confusion matrices do not match!",
+        )
 
 
 class OneHotIoUTest(testing.TestCase):