Axis & Reduction Operations in NumPy

What does axis mean?

Think of axis as the dimension you collapse during a reduction. In a 2D array, axis=0 reduces down the rows (column-wise result), while axis=1 reduces across columns (row-wise result). ND arrays follow the same rule.

import numpy as np

X = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])   # shape (3,3)

print(X.sum(axis=0))  # [12 15 18]  column sums
print(X.sum(axis=1))  # [ 6 15 24]  row sums
print(X.sum())        # 45          all elements

keepdims=True for easy broadcasting

When you plan to subtract or divide a reduced result from the original array, use keepdims=True to preserve singleton dimensions and avoid shape errors.

mu_rows = X.mean(axis=1, keepdims=True)  # shape (3,1)
X_centered_by_row = X - mu_rows          # broadcasts to (3,3)

mu_cols = X.mean(axis=0, keepdims=True)  # shape (1,3)
X_centered_by_col = X - mu_cols

Common reductions

A = np.arange(12).reshape(3,4)

print(A.sum(axis=0))     # column sums
print(A.mean(axis=1))    # row means
print(A.min(axis=0))     # column minima
print(A.max(axis=1))     # row maxima
print(A.std(axis=0))     # column std dev
print(A.prod(axis=1))    # row product

Index of best values: argmax/argmin

To find where the max/min occurs along an axis, use argmax/argmin. Combine with advanced indexing to retrieve the values quickly.

idx = A.argmax(axis=1)          # index of max in each row
row_ids = np.arange(A.shape[0])
row_max_vals = A[row_ids, idx]  # values at those indexes
print(idx, row_max_vals)

Boolean reductions: any / all

Aggregate logical conditions along an axis: “Is there any True?” or “Are all True?”

B = (A % 2 == 0)    # even mask
print(B.any(axis=1))   # any even per row
print(B.all(axis=1))   # all even per row?

Cumulative reductions: cumsum & cumprod

Cumulative versions keep partial results along an axis.

print(A.cumsum(axis=0))  # cumulative column sums
print(A.cumsum(axis=1))  # cumulative row sums

ND examples (3D tensors)

Axis indices start from 0. For (D, H, W), reducing axis=0 collapses depth, axis=1 collapses rows, and axis=2 collapses columns.

T = np.arange(2*3*4).reshape(2,3,4)  # (depth=2, rows=3, cols=4)
print(T.mean(axis=0).shape)  # (3,4)  mean across depth
print(T.sum(axis=2).shape)   # (2,3)  sum across last axis

Typical pitfalls & fixes

• Axis confusion: If shapes don’t match, print .shape and confirm which dimension you reduced.
• Broadcast errors after reduction: Use keepdims=True, or reshape with np.newaxis before arithmetic.
• Integer overflow: For large sums/products on int arrays, consider upcasting to int64 or float64.
• NaNs in data: Use np.nan* variants, e.g., np.nanmean, np.nansum.

Practical patterns

# 1) Standardize columns (Z-score)
X = np.random.randn(5, 3)
mu = X.mean(axis=0, keepdims=True)
sd = X.std(axis=0, keepdims=True)
Z = (X - mu) / sd

# 2) Find the column with max total
col_sums = X.sum(axis=0)
best_col = col_sums.argmax()

# 3) Normalize each row to sum to 1 (avoid div-by-zero)
row_sum = X.sum(axis=1, keepdims=True)
row_sum[row_sum == 0] = 1
X_norm = X / row_sum

Exercises

# a) Given M=(6,4), subtract row-wise medians using keepdims
M = np.arange(24).reshape(6,4)
med = np.median(M, axis=1, keepdims=True)
print(M - med)

# b) From a (3,5) array, return the index of the max per column and the values
A = np.random.randint(0, 100, (3,5))
idx = A.argmax(axis=0)
vals = A[idx, np.arange(A.shape[1])]
print(idx, vals)

# c) For a (2,3,4) tensor, compute cumulative sums along the last axis
T = np.arange(24).reshape(2,3,4)
print(np.cumsum(T, axis=2))

Subhendu Mohapatra

Author & Instructor at plus2net

I write and maintain practical tutorials on Python, PHP, SQL, JavaScript, HTML, jQuery, and web development at plus2net. The tutorials focus on clear explanations, working examples, and code that readers can test and adapt while learning.

About the author   Watch on YouTube

Connect: