Dice loss - incorrect

[td-jakubl]

Tensorflow version 2.16.1 (https://www.tensorflow.org/api_docs/python/tf/keras/losses/Dice)

The calculated dice loss is the average of all values in the tensor, which is not always true e.g. for tensors with shape (batch_size, height, width, 1):

y_true = tf.constant([[[[1.0], [1.0]], [[0.0], [0.0]]], [[[1.0], [1.0]], [[0.0], [0.0]]]])
y_pred = tf.constant([[[[0.0], [1.0]], [[0.0], [1.0]]], [[[0.4], [0.0]], [[0.0], [0.9]]]])

dice = tf.keras.losses.Dice()
print(dice(y_true, y_pred))

The expected result should be tf.Tensor([0.5 0.7575755 ], shape=(2,), dtype=float32) instead of tf.Tensor(0.6164384, shape=(), dtype=float32)

Proposed solution:
Adding an axis parameter for which dimensions the loss is to be calculated, analogous to the: tf.keras.losses.BinaryCrossentropy (https://www.tensorflow.org/api_docs/python/tf/keras/losses/BinaryCrossentropy)

[fchollet]

Sure, we could add axis defaulting to None. Would you be able to open a PR?

[td-jakubl]

Sure, we could add axis defaulting to None. Would you be able to open a PR?

Yes, I am not sure if I created PR correctly, but here it is: #19673

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No