tf.reduce_mean() computes the arithmetic mean of a tensor's elements along one or more specified axes, collapsing that dimension down to the averaged value.

Syntax

tf.reduce_mean(input_tensor, axis=None)

Deep Dive Course

reduce_mean() mirrors reduce_sum() exactly, but computes the average instead of the total — without an axis, it collapses the whole tensor to a single overall mean; with an axis specified, it averages along that dimension only, collapsing it while preserving the others. It's an extremely common operation in machine learning specifically for computing the average loss across a batch of examples, converting a per-example vector of individual loss values into the single scalar loss value that backpropagation actually optimizes against.

1Understanding tf.reduce_mean()

💡

Computing the average loss across a batch with reduce_mean() is one of the most common patterns in a custom training loop — a loss function typically returns one value per example, and reduce_mean() collapses that into the single scalar value needed for computing gradients.

editor.html

import tensorflow as tf

x = tf.constant([1.0, 2.0, 3.0, 4.0])
print(tf.reduce_mean(x))

localhost:3000

2Practical Example

Here is a real-world application of tf.reduce_mean() showing how it is used in production TensorFlow code.

editor.html

import tensorflow as tf

per_example_loss = tf.constant([0.5, 0.3, 0.8, 0.1])
batch_loss = tf.reduce_mean(per_example_loss)
print(batch_loss)

localhost:3000

3Best Practices

Follow these guidelines when working with tf.reduce_mean():

1. Use reduce_mean() to average a per-example loss vector into the single scalar value needed for gradient computation in a custom training loop

2. Specify the axis argument explicitly when you need a per-row or per-column average on multi-dimensional data, rather than collapsing everything into one overall mean

3. Cast integer tensors to float before calling reduce_mean() if an integer division-style truncation isn't the intended behavior

⚠️

Tip: Computing the average loss across a batch with reduce_mean() is one of the most common patterns in a custom training loop — a loss function typically returns one value per example, and reduce_mean() collapses that into the single scalar value needed for computing gradients.

editor.html

import tensorflow as tf

x = tf.constant([1.0, 2.0, 3.0, 4.0])
print(tf.reduce_mean(x))

localhost:3000

Examples

Example 01Basic Usage

import tensorflow as tf

x = tf.constant([1.0, 2.0, 3.0, 4.0])
print(tf.reduce_mean(x))

Example 02Advanced Example

import tensorflow as tf

per_example_loss = tf.constant([0.5, 0.3, 0.8, 0.1])
batch_loss = tf.reduce_mean(per_example_loss)
print(batch_loss)

Best Practices

Use reduce_mean() to average a per-example loss vector into the single scalar value needed for gradient computation in a custom training loop
Specify the axis argument explicitly when you need a per-row or per-column average on multi-dimensional data, rather than collapsing everything into one overall mean
Cast integer tensors to float before calling reduce_mean() if an integer division-style truncation isn't the intended behavior

Interview Question

Why is tf.reduce_mean() typically applied to a loss function's output before computing gradients, rather than backpropagating from a whole vector of per-example losses directly?▼

Hint: Think about what shape a gradient computation ultimately needs to start from.

Backpropagation via automatic differentiation is fundamentally built around computing the gradient of a single scalar value with respect to a set of parameters, since a gradient answers how does this one number change as each parameter changes — a vector of separate per-example losses doesn't have a single unambiguous gradient in that same sense. reduce_mean() collapses the per-example loss vector down into exactly one scalar number representing the batch's overall average loss, which is the single value backpropagation actually needs as its starting point to compute a consistent, well-defined gradient for every trainable parameter in the model.

Exercises

MediumPractice using tf.reduce_mean() in a real scenario.

View Solution

import tensorflow as tf

x = tf.constant([1.0, 2.0, 3.0, 4.0])
print(tf.reduce_mean(x))

Frequently Asked Questions

Why is tf.reduce_mean() typically applied to a loss function's output before computing gradients, rather than backpropagating from a whole vector of per-example losses directly?▼

Related Functions

tf-reduce-sumlosses-meansquarederrornp-mean