sparse_categorical_crossentropy with ignore_class fails for 4D inputs

[enricoschroeder]

Using ignore_class with keras.losses.sparse_categorical_crossentropy and 4D inputs (Batch x Height x Width x Class) fails with a ValueError indicating wrong shapes.

Minimal example to reproduce:

import numpy as np
import tensorflow as tf

y_true = np.zeros((1, 224, 224, 1))
y_true[0, 0, 0, 0] = 255
y_pred = np.ones((1, 224, 224, 21))
tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred, ignore_class=255)

--> "ValueError: Arguments target and output must have the same shape up until the last dimension: target.shape=(1, 224, 224, 1), output.shape=(1, 224, 224, 224, 21)"

This expand_dims seems to be the culprit:

keras/keras/src/losses/losses.py

Line 1719 in 2305fad

ops.expand_dims(valid_mask, -1), y_pred.dtype

[Grvzard]

y_true = np.zeros((1, 224, 224, 1))

=> y_true = np.zeros((1, 224, 224))
Shouldn't y_true has one dimension less than y_pred?

[enricoschroeder]

Oh, you are right, with y_true = np.zeros((1, 224, 224)) it seems to work...

However, when omitting ignore_class from sparse_categorical_crossentropy, y_true = np.zeros((1, 224, 224, 1)) works as well. I was assuming the same behavior regardless of ignore_class... at the very least this should be documented somewhere.