Resolve shape via ops in broadcast_to operation (#547)

keras-team · Jul 19, 2023 · bcb4b0a · bcb4b0a
1 parent 4cfc3eb
commit bcb4b0a
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 9 deletions.
diff --git a/keras_core/metrics/iou_metrics.py b/keras_core/metrics/iou_metrics.py
@@ -106,7 +106,7 @@ def update_state(self, y_true, y_pred, sample_weight=None):
         if len(sample_weight.shape) > 1:
             sample_weight = ops.reshape(sample_weight, [-1])
 
-        sample_weight = ops.broadcast_to(sample_weight, y_true.shape)
+        sample_weight = ops.broadcast_to(sample_weight, ops.shape(y_true))
 
         if self.ignore_class is not None:
             ignore_class = ops.convert_to_tensor(

diff --git a/keras_core/metrics/metric.py b/keras_core/metrics/metric.py
@@ -72,7 +72,9 @@ def update_state(self, y_true, y_pred, sample_weight=None):
             values = ops.cast(values, self.dtype)
             if sample_weight is not None:
                 sample_weight = ops.cast(sample_weight, self.dtype)
-                sample_weight = ops.broadcast_to(sample_weight, values.shape)
+                sample_weight = ops.broadcast_to(
+                    sample_weight, ops.shape(values)
+                )
                 values = ops.multiply(values, sample_weight)
             self.true_positives.assign(self.true_positives + ops.sum(values))
 

diff --git a/keras_core/metrics/metrics_utils.py b/keras_core/metrics/metrics_utils.py
@@ -195,15 +195,15 @@ def _update_confusion_matrix_variables_optimized(
         sample_weights = 1.0
     else:
         sample_weights = ops.broadcast_to(
-            ops.cast(sample_weights, dtype=y_pred.dtype), y_pred.shape
+            ops.cast(sample_weights, dtype=y_pred.dtype), ops.shape(y_pred)
         )
         if not multi_label:
             sample_weights = ops.reshape(sample_weights, [-1])
     if label_weights is None:
         label_weights = 1.0
     else:
         label_weights = ops.expand_dims(label_weights, 0)
-        label_weights = ops.broadcast_to(label_weights, y_pred.shape)
+        label_weights = ops.broadcast_to(label_weights, ops.shape(y_pred))
         if not multi_label:
             label_weights = ops.reshape(label_weights, [-1])
     weights = ops.cast(
@@ -533,7 +533,7 @@ def update_confusion_matrix_variables(
 
     if label_weights is not None and not multi_label:
         label_weights = ops.expand_dims(label_weights, 0)
-        label_weights = ops.broadcast_to(label_weights, y_pred.shape)
+        label_weights = ops.broadcast_to(label_weights, ops.shape(y_pred))
         label_weights_tiled = ops.tile(
             ops.reshape(label_weights, thresh_tiles), data_tiles
         )

diff --git a/keras_core/metrics/regression_metrics.py b/keras_core/metrics/regression_metrics.py
@@ -428,7 +428,6 @@ def __init__(
             shape=(),
             initializer=initializers.Zeros(),
             name="num_samples",
-            dtype="int32",
         )
         self._built = False
 
@@ -500,16 +499,19 @@ def update_state(self, y_true, y_pred, sample_weight=None):
             # Make sure there's a features dimension
             sample_weight = ops.expand_dims(sample_weight, axis=1)
 
-        sample_weight = ops.broadcast_to(sample_weight, y_true.shape)
+        sample_weight = ops.broadcast_to(sample_weight, ops.shape(y_true))
 
-        weighted_y_true = y_true * sample_weight
+        weighted_y_true = y_true * ops.cast(sample_weight, y_true.dtype)
         self.sum.assign(self.sum + ops.sum(weighted_y_true, axis=0))
         self.squared_sum.assign(
             self.squared_sum + ops.sum(y_true * weighted_y_true, axis=0)
         )
         self.total_mse.assign(
             self.total_mse
-            + ops.sum((y_true - y_pred) ** 2 * sample_weight, axis=0)
+            + ops.sum(
+                (y_true - y_pred) ** 2 * ops.cast(sample_weight, y_true.dtype),
+                axis=0,
+            )
         )
         self.count.assign(self.count + ops.sum(sample_weight, axis=0))
         self.num_samples.assign(self.num_samples + ops.size(y_true))