fix: validate sparse_categorical_crossentropy class indices off CPU for Torch backend - #23669

This PR fixes , keras.ops.sparse_categorical_crossentropy only checks the shape of the inputs. It does not check that the label values are valid. It passes them straight to F.cross_entropy / F.nll_loss on the Torch backend.

Those PyTorch functions only check label values on CPU. So the same bad label behaves differently on each device:

Example with 3 classes, where label 5 is invalid:

output = torch.tensor([[0.9, 0.05, 0.05], [0.1, 0.8, 0.1]])  # 3 classes
target = torch.tensor([0, 5])                                # 5 is invalid

with keras.device("cpu"):
    ops.sparse_categorical_crossentropy(target, output)
    # IndexError: Target 5 is out of bounds.

with keras.device("mps"):
    ops.sparse_categorical_crossentropy(target, output)
    # tensor([0.1054, -0.0000], device='mps:0')   <- no error, wrong loss

This is a real problem. If your labels and your number of classes do not match, training on MPS keeps running and slowly learns the wrong thing. Nothing tells you something is wrong.

The fix

On non-CPU devices, bad labels are changed to 0 so the lookup stays in range. Then the loss for those rows is set to NaN:

def _mask_invalid_class_indices(target, num_classes):
    if target.device.type == "cpu":
        return target, None  # torch already raises here
    invalid = ~((target == -100) | ((target >= 0) & (target < num_classes)))
    return target.masked_fill(invalid, 0), invalid

    # The kernels below only validate class indices on CPU.
    target, invalid = _mask_invalid_class_indices(target, output.shape[1])
    ...
    if invalid is not None:
        result = result.masked_fill(invalid, float("nan"))

This runs after the class axis is moved to position 1. So output.shape[1] is always the number of classes, no matter the rank or the axis value. One call covers both the logits path and the probabilities path.

Result for target=[0, 5] with 3 classes:

Why not just raise an error everywhere

To raise an error, we have to copy the labels from the GPU back to the CPU and look at them. That makes the GPU wait, on every training step.

I measured this on an M4 with MPS, batch=4096, vocab=32000:

The new check costs nothing. The error version does cost something, on every step.

The check also has no if statement that depends on the data. That means torch.compile can trace it with fullgraph=True, and it also works on meta tensors during shape inference.

One small difference from TensorFlow

TensorFlow makes both the loss and the gradients NaN. Here, only the loss is NaN. masked_fill blocks the gradient, so the gradient for a bad row is 0.

This means a bad label cannot push the weights toward the wrong class, and it cannot turn the weights into NaN. You see nan in the loss right away, but the model is still usable. This is written in the docstring and tested by test_invalid_target_contributes_no_gradient.

Testing

Tested on an Apple M4 with a real MPS device, torch 2.12.1:

• The bug happens on master and is fixed on this branch.
• It works through the public API too: keras.losses.SparseCategoricalCrossentropy() returns nan, and model.fit shows loss: nan instead of a normal-looking number.
• keras/src/backend/torch/nn_test.py — 30 passed (12 on real MPS).
• Same tests with MPS and CUDA turned off, to copy CI — 13 passed, 1 skipped.
• keras/src/losses/losses_test.py -k sparse — 19 passed, nothing broken.
• ruff check and ruff format are clean.

Fixes: #23664

Contributor Agreement

Please review our AI-Assisted Contribution Policy and check all boxes below before submitting your PR for review:

• I am a human, and not a bot.
• I will be responsible for responding to review comments in a timely manner.
• I will work with the maintainers to push this PR forward until submission.

Note: Failing to adhere to this agreement may result in your future PRs no longer being reviewed.