Torch backend: `binary_crossentropy` gives a finite loss for NaN predictions on MPS

Description

On the Torch backend with the MPS device, keras.ops.binary_crossentropy gives a loss of 100.0 for a NaN prediction when from_logits=False. The JAX, NumPy, and TensorFlow backends give NaN. Torch on CPU raises an error.

A model with NaN weights gets a finite loss on MPS. keras.callbacks.TerminateOnNaN does not find a NaN loss, and the training does not stop.

Procedure

Run this code on an Apple silicon Mac:

import os

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

import keras

target = keras.ops.convert_to_tensor([1.0, 1.0, 0.0, 0.0])
output = keras.ops.convert_to_tensor([0.3, float("nan"), 0.3, 0.9])
print(keras.ops.binary_crossentropy(target, output))

This second example shows the effect on TerminateOnNaN:

import os

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

import numpy as np

import keras

model = keras.Sequential(
    [keras.Input((4,)), keras.layers.Dense(1, activation="sigmoid")]
)
kernel, bias = model.get_weights()
kernel[0, 0] = np.nan
model.set_weights([kernel, bias])
model.compile(optimizer="sgd", loss="binary_crossentropy")

x = np.random.rand(64, 4).astype("float32")
y = np.random.randint(0, 2, (64, 1)).astype("float32")
history = model.fit(
    x,
    y,
    epochs=3,
    batch_size=32,
    verbose=0,
    callbacks=[keras.callbacks.TerminateOnNaN()],
)
print(history.history["loss"])

Actual result

On MPS, the first example prints:

tensor([  1.2040, 100.0000,   0.3567,   2.3026], device='mps:0')

The second example completes all 3 epochs and prints [100.0, 100.0, 100.0].

Expected result

The loss for a NaN prediction is NaN, as on the other backends. TerminateOnNaN stops the training after the first batch.

Results on other backends and devices

Backend	Device	First example	Second example
JAX	CPU	[1.2040, nan, 0.3567, 2.3026]	Stops after batch 0. The loss is [nan].
NumPy	CPU	[1.2040, nan, 0.3567, 2.3026]	Not applicable
TensorFlow	CPU	[1.2040, nan, 0.3567, 2.3026]	Stops after batch 0. The loss is [nan].
Torch	CPU	Raises RuntimeError: all elements of input should be between 0 and 1	Raises the same RuntimeError
Torch	MPS	[1.2040, 100.0, 0.3567, 2.3026]	Completes 3 epochs. The loss is [100.0, 100.0, 100.0].

With from_logits=True, MPS gives NaN for the NaN prediction. Only from_logits=False has the problem.

Cause

binary_crossentropy in keras/src/backend/torch/nn.py clips output with torch.clip. The clip keeps the NaN values. The function then calls torch.nn.functional.binary_cross_entropy. The MPS kernel computes max(log(x), -100) and max(log(1 - x), -100). MPSGraph gives -100 for these terms when x is NaN. For this reason, the loss is 100.0. The CPU kernel checks the input range and raises an error.

The same result occurs without Keras:

import torch
import torch.nn.functional as F

output = torch.tensor([0.3, float("nan")], device="mps")
target = torch.tensor([1.0, 1.0], device="mps")
print(F.binary_cross_entropy(output, target, reduction="none"))
# tensor([  1.2040, 100.0000], device='mps:0')

PR #23633 changes this function. The problem also occurs with the change from that PR.

Possible fix

Keep the NaN values in the result. For example, replace NaN values in the clipped output before the kernel call. Then use torch.where(torch.isnan(output), output, loss) on the result. With this change, Torch on MPS and on CPU gives NaN, as the other backends do.

Environment

• Keras: master at 32530b6ed
• PyTorch: 2.10.0
• Python: 3.12.12
• macOS 26.6.2 on an Apple M1 Pro