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

🎥 Join me live on YouTube

Passionate about coding and teaching, I publish practical tutorials on PHP, Python, JavaScript, SQL, and web development. My goal is to make learning simple, engaging, and project‑oriented with real examples and source code.

← Subscribe to our YouTube Channel here