#23621·keras

Inconsistent behavior of binary_crossentropy under torch backend

Author: stefanradev93Created Sep 13, 2026Updated Sep 15, 2026
Labelstype:Bugbackend:torch

Description

The Torch backend of keras.ops.binary_crossentropy squeezes the final dimension of targets and predictions when MPS is available and that dimension has size 1. It returns the loss without restoring the original shape.

Consequently, identical inputs produce different output shapes depending on MPS availability:

  • CPU/CUDA without MPS: (batch, steps, 1)
  • MPS available: (batch, steps)

This breaks code that expects elementwise binary cross-entropy to preserve the input shape.

The condition also checks torch.backends.mps.is_available() rather than the tensors' actual device, so it affects CPU tensors on machines where MPS is available.

Reproduction

Run on a Mac with MPS available:

import os
os.environ["KERAS_BACKEND"] = "torch"

import keras
import torch

assert torch.backends.mps.is_available()

with keras.device("cpu"):
    targets = keras.ops.zeros((3, 4, 1))
    logits = keras.ops.zeros((3, 4, 1))
    loss = keras.ops.binary_crossentropy(
        targets, logits, from_logits=True
    )

    print("Input shape:", targets.shape)
    print("Loss shape:", loss.shape)
    assert tuple(loss.shape) == (3, 4, 1)

Expected behavior

The loss has shape (3, 4, 1) consistently across devices.

Actual behavior

When MPS is available, the loss has shape (3, 4). For example, downstream code using loss[..., 0] then produces (3,) instead of (3, 4), causes shape mismatches with e.g., sequence masks.

Relevant implementation

In keras/src/backend/torch/nn.py, binary_crossentropy explicitly squeezes both inputs before computing the loss, without restoring the dimension afterward. This implementation is present in Keras 3.15.1.

Suggested fix

Restore the original shape before returning the loss. If the internal squeeze is required only for MPS tensors, derive it from the tensors' actual device rather than global MPS availability.