`sparse_categorical_crossentropy` ignores feature channel (and axis parameter) when checking tensor dimensions [torch backend])

[enricoschroeder]

I have two tensors which I want to apply sparse_categorical_crossentropy to.

I am using the torch backend, so the tensors are of shape (batch, features, x, y).

However, the function ignores the fact that the feature dimension with torch backend is not the last dimension but the first (non-batch) dimension and thus throws ValueError: Arguments "target" and "output" must have the same shape up until the last dimension: target.shape=torch.Size([16, 1, 224, 224]), output.shape=torch.Size([16, 2, 224, 224]).

This happens even when specifying axis=1 (which by the way is not exposed in keras.losses.SparseCategoricalCrossentropy and thus not an option when using the loss with model.fit).

sparse_categorical_crossentropy should respect the feature channel setting or at least the axis parameter.

[sonali-kumari1]

Hi @enricoschroeder -
Could you please provide a code snippet to reproduce the reported behavior? Thanks!

[enricoschroeder]

Sure, below is a reproducible example. This is with torch backend and "image_data_format" = "channels_first".

import keras
import numpy as np

y_true = np.ones([16, 1, 224, 224], dtype="float32")
y_pred = np.ones([16, 2, 224, 224], dtype="float32")
y_pred[:, 0, :, :] = 0

keras.losses.sparse_categorical_crossentropy(y_true, y_pred, axis=1)

Expected output: Loss value for the two tensors.
Actual output: ValueError: Arguments "target" and "output" must have the same shape up until the last dimension: target.shape=torch.Size([16, 1, 224, 224]), output.shape=torch.Size([16, 2, 224, 224])

[edge7]

This should be an easy fix. The main issue is that axis is not propagated to the keras ops (which does offer that parameter in the signature)

[sonali-kumari1]

Hi @enricoschroeder, @edge7 -
I tried to reproduce this issue with the latest version of keras(3.9.0) across Pytorch, Tensorflow and Jax backends and encountered the following errors:

• Pytorch
  ValueError: Arguments targetandoutput must have the same shape up until the last dimension: target.shape=torch.Size([16, 1, 224, 224]), output.shape=torch.Size([16, 2, 224, 224])
• Jax
  ValueError: Arguments targetandoutput must have the same shape up until the last dimension: target.shape=(16, 1, 224, 224), output.shape=(16, 2, 224, 224)
• Tensorflow
  ValueError: Argument output must have rank (ndim) target.ndim - 1. Received: target.shape=(16, 1, 224, 224), output.shape=(16, 2, 224, 224)

I am attaching gist for your reference.

[edge7]

I opened a PR, if merged one could pass the axis parameter down to the real function (Torch only).
The only thing is that the target (the ground truth) should already be squeezed