Introduction

In Keras, a custom loss does not usually need to collapse the entire batch into one number by hand. The normal contract is that the loss computes a value for each sample, and Keras then applies its configured reduction to turn those values into the scalar objective used for backpropagation.

The Short Answer

When you pass a loss to model.compile, Keras is generally happy with a tensor of per-sample losses. For a batch of size N, that often means a result shaped like (N,) after reducing any feature dimension inside each sample.

Keras then reduces those per-sample losses according to the selected reduction behavior. In the common case, that becomes a single scalar such as the mean batch loss.

So the practical rule is:

• inside the loss, reduce over the prediction dimensions that belong to one sample
• let Keras reduce across the batch unless you have a specific reason not to

What Built-In Losses Do

A built-in loss such as mean squared error illustrates the pattern:

1import tensorflow as tf
2
3y_true = tf.constant([[1.0, 0.0], [0.0, 1.0]])
4y_pred = tf.constant([[0.8, 0.2], [0.3, 0.7]])
5
6per_sample = tf.reduce_mean(tf.square(y_true - y_pred), axis=-1)
7print(per_sample.numpy())

Output:

[0.04 0.09]

The loss above returns one value per row in the batch. Keras can then average those values during training.

If you write the loss as a Loss class, Keras still expects the call method to produce unreduced or partially reduced loss values that it can aggregate according to the configured reduction.

A Correct Custom Loss

Here is a custom loss function that returns one value per sample:

1import tensorflow as tf
2from tensorflow import keras
3
4def clipped_mae(y_true, y_pred):
5    error = tf.abs(y_true - y_pred)
6    error = tf.minimum(error, 2.0)
7    return tf.reduce_mean(error, axis=-1)
8
9model = keras.Sequential([
10    keras.layers.Input(shape=(3,)),
11    keras.layers.Dense(4, activation="relu"),
12    keras.layers.Dense(1)
13])
14
15model.compile(optimizer="adam", loss=clipped_mae)

axis=-1 removes the feature dimension for each sample. The batch dimension remains, which is what Keras wants in most cases.

When a Scalar Is Also Acceptable

You can return a scalar for the whole batch if you intentionally reduce everything yourself. TensorFlow can still differentiate that scalar. However, doing so means you take control over batch reduction semantics.

1import tensorflow as tf
2
3def batch_loss(y_true, y_pred):
4    error = tf.square(y_true - y_pred)
5    return tf.reduce_mean(error)

This works, but it is usually less flexible:

• sample weighting becomes harder to reason about
• distributed training behavior is easier to get wrong
• you may accidentally double-reduce if you mix APIs carelessly

That is why per-sample loss values are the safer default.

Why Shape Matters

Keras distinguishes between sample dimensions and feature dimensions. Suppose your model outputs shape (batch_size, 10) for ten regression targets. The loss should usually reduce across the 10 outputs for each sample, not across the entire batch.

def per_sample_mse(y_true, y_pred):
    squared = tf.square(y_true - y_pred)
    return tf.reduce_mean(squared, axis=-1)

If you forget axis=-1, the loss tensor may keep extra dimensions such as (batch_size, 10). Sometimes Keras can still reduce it, but the result may not match your intention. Always think explicitly about which axes represent one example and which axis represents the batch.

Using a Loss Class

If you need configuration, subclass keras.losses.Loss:

1import tensorflow as tf
2from tensorflow import keras
3
4class HuberLikeLoss(keras.losses.Loss):
5    def __init__(self, delta=1.0, name="huber_like"):
6        super().__init__(name=name)
7        self.delta = delta
8
9    def call(self, y_true, y_pred):
10        error = y_true - y_pred
11        abs_error = tf.abs(error)
12        quadratic = tf.minimum(abs_error, self.delta)
13        linear = abs_error - quadratic
14        loss = 0.5 * tf.square(quadratic) + self.delta * linear
15        return tf.reduce_mean(loss, axis=-1)

This preserves the standard Keras pattern: compute a per-sample loss, then let the framework reduce across the batch.

Mental Model for Backpropagation

Backpropagation still needs one scalar objective in the end. The question is not whether a scalar is required eventually; it is where that scalar should be formed.

In Keras, the clean answer is usually:

1. your loss computes per-sample values
2. Keras applies sample weights and reduction
3. the training loop obtains the final scalar used for gradients

That division of responsibility keeps custom losses composable with the rest of the framework.

Common Pitfalls

• Returning a tensor with extra feature dimensions because you forgot to reduce across axis=-1. That often produces unexpected training behavior.
• Reducing over the whole batch inside the loss without meaning to. This can interfere with weighting and distributed execution semantics.
• Assuming the loss must always return exactly one scalar from the function. In standard Keras usage, per-sample loss values are normal.
• Writing a mathematically correct loss that uses Python branching instead of TensorFlow ops. That can break graph execution and differentiation.
• Ignoring sample weights and masking when designing the loss. Letting Keras handle batch reduction makes those features easier to preserve.

Summary

• Keras custom losses typically return one value per sample, not one manually computed scalar for the whole batch.
• Keras then reduces those values into the scalar objective needed for gradient updates.
• Returning a single batch scalar can work, but it is usually less flexible.
• The most important design choice is reducing over feature axes while keeping the batch axis intact.
• If you match Keras' expected shapes, custom losses stay compatible with weighting, masking, and distribution strategies.