Introduction

TensorFlow has built-in reductions such as mean, sum, and max, but median is less direct. You usually compute it by sorting values and selecting the middle element, or by using TensorFlow Probability if that dependency is available. This article covers both approaches and shows how to handle odd-sized, even-sized, and axis-based medians.

Median for a One-Dimensional Tensor

The median is the middle value after sorting. For an odd number of items, it is the center element. For an even number, it is usually the average of the two center elements.

Here is a reusable TensorFlow implementation for a one-dimensional tensor:

1import tensorflow as tf
2
3
4def median_1d(x: tf.Tensor) -> tf.Tensor:
5    x = tf.convert_to_tensor(x)
6    x = tf.reshape(x, [-1])
7    x = tf.sort(x)
8
9    n = tf.shape(x)[0]
10    mid = n // 2
11
12    def odd():
13        return x[mid]
14
15    def even():
16        left = tf.cast(x[mid - 1], tf.float32)
17        right = tf.cast(x[mid], tf.float32)
18        return (left + right) / 2.0
19
20    return tf.cond(tf.equal(n % 2, 1), odd, even)
21
22
23print(median_1d(tf.constant([7, 1, 3, 9, 5])).numpy())   # 5
24print(median_1d(tf.constant([7, 1, 3, 9])).numpy())      # 5.0

This approach is simple and works in eager mode and inside tf.function.

Why tf.reduce_mean Is Not Enough

Median and mean answer different questions. Mean is influenced strongly by outliers, while median is robust when the distribution contains extreme values.

1import tensorflow as tf
2
3x = tf.constant([1.0, 2.0, 3.0, 1000.0])
4print(tf.reduce_mean(x).numpy())  # 251.5
5print(median_1d(x).numpy())       # 2.5

For noisy or skewed data, the median often reflects the typical value more faithfully than the mean.

Median Along an Axis

For matrices or higher-rank tensors, you often want the median per row or per column. One practical solution is to sort along the chosen axis and select the middle positions.

1import tensorflow as tf
2
3
4def median_axis_0(x: tf.Tensor) -> tf.Tensor:
5    x = tf.convert_to_tensor(x)
6    x = tf.sort(x, axis=0)
7    n = tf.shape(x)[0]
8    mid = n // 2
9
10    def odd():
11        return x[mid]
12
13    def even():
14        left = tf.cast(x[mid - 1], tf.float32)
15        right = tf.cast(x[mid], tf.float32)
16        return (left + right) / 2.0
17
18    return tf.cond(tf.equal(n % 2, 1), odd, even)
19
20
21matrix = tf.constant([
22    [1.0, 10.0],
23    [3.0, 30.0],
24    [2.0, 20.0],
25    [4.0, 40.0],
26])
27
28print(median_axis_0(matrix).numpy())  # [ 2.5 25. ]

The same idea can be adapted for other axes by transposing or sorting on the axis you care about.

Using TensorFlow Probability

If you already depend on TensorFlow Probability, tfp.stats.percentile is often the cleanest option because the median is simply the 50th percentile.

1import tensorflow as tf
2import tensorflow_probability as tfp
3
4x = tf.constant([7.0, 1.0, 3.0, 9.0, 5.0])
5median = tfp.stats.percentile(x, 50.0, interpolation="midpoint")
6print(median.numpy())

This can be especially convenient when your project already computes other percentiles.

Performance Considerations

Median usually requires sorting, which is more expensive than reductions like sum or mean. For very large tensors, that extra cost matters.

Practical guidance:

• use median only where robustness to outliers matters
• avoid recomputing it repeatedly if data is unchanged
• compute along the smallest useful axis

If you only need a rough robust statistic, a percentile approximation outside TensorFlow may sometimes be more efficient.

Dtype Details

For even-sized inputs, averaging the middle pair can promote integer data to floating-point output. That is usually what you want, because the true median may lie between two integers.

If you require integer output for a specific workflow, define the rounding behavior explicitly instead of relying on implicit casts.

Median in Model Pipelines

Median can be useful in TensorFlow pipelines for:

• robust feature summaries
• filtering noisy signals
• diagnostics during preprocessing

It is less common inside the core forward pass of neural networks, but it appears frequently in data-cleaning and analysis workflows around training.

Common Pitfalls

• Assuming TensorFlow has a direct tf.reduce_median in core APIs.
• Forgetting to handle even-length inputs separately from odd-length inputs.
• Returning integer output when the correct median should be fractional.
• Sorting along the wrong axis and getting a valid but incorrect result.
• Using median in hot paths without considering the cost of sorting.

Summary

• TensorFlow median is usually implemented by sorting and selecting the middle values.
• Odd-sized and even-sized tensors need slightly different handling.
• TensorFlow Probability can compute median through percentile utilities.
• Median is more robust than mean when data contains outliers.
• Pay attention to axis choice, dtype conversion, and sorting cost in production code.