Resolve shape via ops in broadcast_to operation (#547)

This commit is contained in:
Ramesh Sampath 2023-07-19 15:04:11 -05:00 committed by Francois Chollet
parent 2c2962158e
commit 42bdabf76a
4 changed files with 13 additions and 9 deletions

@ -106,7 +106,7 @@ class _IoUBase(Metric):
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(

@ -72,7 +72,9 @@ class Metric:
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))

@ -195,7 +195,7 @@ 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])
@ -203,7 +203,7 @@ def _update_confusion_matrix_variables_optimized(
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
)

@ -428,7 +428,6 @@ class R2Score(reduction_metrics.Metric):
shape=(),
initializer=initializers.Zeros(),
name="num_samples",
dtype="int32",
)
self._built = False
@ -500,16 +499,19 @@ class R2Score(reduction_metrics.Metric):
# 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))