NumPy Axis and Reduction: Predict Shapes and Verify Results

An axis is one dimension of an array. A reduction combines values along that dimension; the remaining dimensions describe the result. Start with the input shape and the meaning of each dimension, then choose the axis and statistic. Use keepdims when the result must broadcast back to the input.

What Does axis Mean? 🔝

An axis is an array dimension. A reduction combines values along the chosen dimension, usually removing that dimension from the result. In a 2D array, axis=0 reduces down the rows, leaving column results; axis=1 reduces across columns, leaving row results. Arrays do not carry business labels automatically.

import numpy as np

def sample_array():
    return np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

X = sample_array()
print(X)
print("Column sums:", X.sum(axis=0))
print("Row sums:", X.sum(axis=1))
print("All elements:", X.sum())
np.testing.assert_array_equal(X.sum(axis=0), [12, 15, 18])
np.testing.assert_array_equal(X.sum(axis=1), [6, 15, 24])
assert X.sum() == 45

Expected output

[[1 2 3]
 [4 5 6]
 [7 8 9]]
Column sums: [12 15 18]
Row sums: [ 6 15 24]
All elements: 45

Predict Shapes with a Rectangular Array 🔝

A square array gives row and column reductions the same shape, which can hide a mistaken axis. Use a rectangular example to make the difference visible. For shape (3, 4), reducing axis=0 leaves (4,), while reducing axis=1 leaves (3,). See array shapes.

A = np.arange(12).reshape(3, 4)
print("Input:", A.shape)
print("axis=0:", A.sum(axis=0).shape)
print("axis=1:", A.sum(axis=1).shape)
print("axis=None:", np.shape(A.sum()))
assert A.sum(axis=0).shape == (4,)
assert A.sum(axis=1).shape == (3,)
assert np.shape(A.sum()) == ()

Expected output

Input: (3, 4)
axis=0: (4,)
axis=1: (3,)
axis=None: ()

keepdims for Row and Column Centering 🔝

keepdims retains reduced axes with size one. Row means shaped (3, 1) broadcast across columns; column means shaped (1, 3) broadcast down rows. Subtracting either mean produces a different result. Link to broadcasting for the general compatibility rule.

X = sample_array()
mu_rows = X.mean(axis=1, keepdims=True)
mu_cols = X.mean(axis=0, keepdims=True)
rows_centered = X - mu_rows
cols_centered = X - mu_cols
print("Row-mean shape:", mu_rows.shape)
print("Row-centered:")
print(rows_centered)
print("Column-mean shape:", mu_cols.shape)
print("Column-centered:")
print(cols_centered)
np.testing.assert_allclose(rows_centered.mean(axis=1), 0)
np.testing.assert_allclose(cols_centered.mean(axis=0), 0)

Expected output

Row-mean shape: (3, 1)
Row-centered:
[[-1.  0.  1.]
 [-1.  0.  1.]
 [-1.  0.  1.]]
Column-mean shape: (1, 3)
Column-centered:
[[-3. -3. -3.]
 [ 0.  0.  0.]
 [ 3.  3.  3.]]

A Broadcasting Mistake Can Run without an Error 🔝

On a square array, subtracting a (3,) row-mean vector aligns it with the last dimension, so the calculation runs but subtracts the wrong values. Retain the row axis correctly with keepdims or add a singleton dimension using [:, None]. Shape compatibility alone does not prove the intended calculation.

X = sample_array()
row_means = X.mean(axis=1)
wrong = X - row_means
correct = X - row_means[:, None]
print("Wrong row means after subtraction:", wrong.mean(axis=1))
print("Correct row means after subtraction:", correct.mean(axis=1))
assert not np.allclose(wrong.mean(axis=1), 0)
np.testing.assert_allclose(correct.mean(axis=1), 0)

Expected output

Wrong row means after subtraction: [-3.  0.  3.]
Correct row means after subtraction: [0. 0. 0.]

Common Reductions: sum, mean, min, max, std and prod 🔝

The axis meaning stays the same across these functions, but the measure changes. std defaults to ddof=0, the population standard deviation. Choose the statistic and its assumptions deliberately rather than only matching an output shape.

A = np.arange(12).reshape(3, 4)
print("Column sums:", A.sum(axis=0))
print("Row means:", A.mean(axis=1))
print("Column minima:", A.min(axis=0))
print("Row maxima:", A.max(axis=1))
print("Column population std:", A.std(axis=0).round(3))
print("Row products:", A.prod(axis=1))
np.testing.assert_array_equal(A.sum(axis=0), [12, 15, 18, 21])
np.testing.assert_allclose(A.mean(axis=1), [1.5, 5.5, 9.5])

Expected output

Column sums: [12 15 18 21]
Row means: [1.5 5.5 9.5]
Column minima: [0 1 2 3]
Row maxima: [ 3  7 11]
Column population std: [3.266 3.266 3.266 3.266]
Row products: [   0  840 7920]

argmax and argmin Return Positions 🔝

argmax returns the first position of the maximum along the selected axis, not the value or every tie. Use the positions to retrieve values. A global argmax is a flattened position; unravel_index maps it back to coordinates. Guard empty and invalid data before selecting an extreme.

A = np.arange(12).reshape(3, 4)
positions = A.argmax(axis=1)
row_ids = np.arange(A.shape[0])
values = A[row_ids, positions]
print("Maximum positions per row:", positions)
print("Maximum values:", values)
print("Minimum positions per row:", A.argmin(axis=1))
print("Global maximum coordinates:", np.unravel_index(A.argmax(), A.shape))
np.testing.assert_array_equal(values, A.max(axis=1))
ties = np.array([[7, 7, 2]])
assert ties.argmax(axis=1)[0] == 0
print("All tied positions:", np.flatnonzero(ties[0] == ties[0].max()))

Expected output

Maximum positions per row: [3 3 3]
Maximum values: [ 3  7 11]
Minimum positions per row: [0 0 0]
Global maximum coordinates: (np.int64(2), np.int64(3))
All tied positions: [0 1]

any and all Reduce Boolean Conditions 🔝

any asks whether at least one selected condition is True; all asks whether every selected condition is True. Summing a Boolean mask counts matches. For a row-validation rule, also account for missing measurements instead of treating comparisons with NaN as ordinary observed values.

A = np.arange(12).reshape(3, 4)
even = A % 2 == 0
print("Any even per row:", even.any(axis=1))
print("All even per row:", even.all(axis=1))
print("Even counts per row:", even.sum(axis=1))
np.testing.assert_array_equal(even.any(axis=1), [True, True, True])
np.testing.assert_array_equal(even.all(axis=1), [False, False, False])
np.testing.assert_array_equal(even.sum(axis=1), [2, 2, 2])

Expected output

Any even per row: [ True  True  True]
All even per row: [False False False]
Even counts per row: [2 2 2]

cumsum and cumprod Keep Running Results 🔝

A cumulative operation retains partial results along an axis rather than collapsing it. With axis=None, cumsum flattens the array. Specify the axis to keep the day/product layout meaningful. Products and sums can overflow narrow integer types.

A = np.arange(12).reshape(3, 4)
print("Cumulative columns:")
print(A.cumsum(axis=0))
print("Cumulative rows:")
print(A.cumsum(axis=1))
factors = np.array([[2, 3, 4], [1, 2, 3]])
print("Cumulative row products:")
print(factors.cumprod(axis=1))
assert A.cumsum(axis=0).shape == A.shape
np.testing.assert_array_equal(A.cumsum(axis=1)[:, -1], A.sum(axis=1))

Expected output

Cumulative columns:
[[ 0  1  2  3]
 [ 4  6  8 10]
 [12 15 18 21]]
Cumulative rows:
[[ 0  1  3  6]
 [ 4  9 15 22]
 [ 8 17 27 38]]
Cumulative row products:
[[ 2  6 24]
 [ 1  2  6]]

Three-Dimensional Arrays and Several Axes 🔝

For shape (2, 3, 4), axis=0 is the first dimension and axis=2 is the last. Reducing a tuple removes several dimensions at once. Negative axis numbers count from the end; axis=-1 is the last axis. Predict which dimensions remain before choosing the reduction.

T = np.arange(24).reshape(2, 3, 4)
print("Input:", T.shape)
print("Mean axis=0:", T.mean(axis=0).shape)
print("Sum axis=2:", T.sum(axis=2).shape)
print("Sum axes=(1, 2):", T.sum(axis=(1, 2)))
print("Retained dimensions:", T.sum(axis=(1, 2), keepdims=True).shape)
np.testing.assert_array_equal(T.sum(axis=-1), T.sum(axis=2))
np.testing.assert_array_equal(T.sum(axis=(1, 2)), [66, 210])

Expected output

Input: (2, 3, 4)
Mean axis=0: (3, 4)
Sum axis=2: (2, 3)
Sum axes=(1, 2): [ 66 210]
Retained dimensions: (2, 1, 1)

Missing Data: Subtotals and Valid Counts 🔝

A NaN propagates through ordinary mean or sum. NaN-aware functions ignore missing values, but counts still matter: an all-NaN nansum is zero and an all-NaN mean has no known result. Calculate an available-value mean with guarded division to keep an all-missing row unknown without a divide-by-zero warning.

values = np.array([[10.0, np.nan, 20.0], [np.nan, np.nan, np.nan]])
valid = ~np.isnan(values)
totals = np.sum(values, axis=1, where=valid)
counts = valid.sum(axis=1)
means = np.divide(totals, counts, out=np.full(totals.shape, np.nan), where=counts > 0)
print("Known subtotals:", totals)
print("Known counts:", counts)
print("Available-value means:", means)
assert means[0] == 15 and np.isnan(means[1])
assert np.isnan(values.sum(axis=1)).all()

Expected output

Known subtotals: [30.  0.]
Known counts: [2 0]
Available-value means: [15. nan]

Standardize Columns and Handle Constant Features 🔝

Column standardization subtracts each feature mean and divides by its standard deviation. Here one feature is constant, so its standard deviation is zero. The chosen policy encodes that feature as zero and reports the constant-feature mask. Other applications may drop it. Fit preprocessing on training data only when building a predictive model.

X = np.array([[1.0, 10.0, 5.0], [2.0, 20.0, 5.0],
              [3.0, 30.0, 5.0], [4.0, 40.0, 5.0], [5.0, 50.0, 5.0]])
mu = X.mean(axis=0, keepdims=True)
sd = X.std(axis=0, keepdims=True)
Z = np.divide(X - mu, sd, out=np.zeros_like(X), where=sd > 0)
print("Constant feature mask:", sd == 0)
print("Standardized features:")
print(Z.round(3))
np.testing.assert_allclose(Z.mean(axis=0), 0, atol=1e-12)
np.testing.assert_allclose(Z.std(axis=0)[:,], [1, 1, 0])

Expected output

Constant feature mask: [[False False  True]]
Standardized features:
[[-1.414 -1.414  0.   ]
 [-0.707 -0.707  0.   ]
 [ 0.     0.     0.   ]
 [ 0.707  0.707  0.   ]
 [ 1.414  1.414  0.   ]]

Row Shares Need a Meaningful Denominator 🔝

Use nonnegative quantities for within-row shares. An all-zero row has no defined share denominator. This example retains NaN for that row instead of claiming that its shares add to one. This interpretation is not appropriate for arbitrary signed, standardized values.

units = np.array([[2.0, 3.0, 5.0], [0.0, 0.0, 0.0], [1.0, 1.0, 2.0]])
assert np.isfinite(units).all() and (units >= 0).all()
totals = units.sum(axis=1, keepdims=True)
shares = np.divide(units, totals, out=np.full_like(units, np.nan), where=totals > 0)
print("Row shares:")
print(shares)
np.testing.assert_allclose(shares[[0, 2]].sum(axis=1), [1, 1])
assert np.isnan(shares[1]).all()

Expected output

Row shares:
[[0.2  0.3  0.5 ]
 [ nan  nan  nan]
 [0.25 0.25 0.5 ]]

Practical Sales Matrix: Days versus Products 🔝

This synthetic matrix contains three days and four products. axis=0 answers which product sold the most units across days; axis=1 answers how many units were sold each day. These are different business questions, even though both use sum. Retain the labels outside the numerical array.

products = np.array(["Books", "Pens", "Folders", "Gifts"])
sales = np.array([[12, 6, 3, 2], [8, 10, 4, 1], [10, 4, 5, 3]])
product_totals = sales.sum(axis=0, dtype=np.int64)
daily_totals = sales.sum(axis=1, dtype=np.int64)
print("Product totals:", dict(zip(products.tolist(), product_totals.tolist())))
print("Daily totals:", daily_totals)
print("Highest-total product:", products[product_totals.argmax()])
np.testing.assert_array_equal(product_totals, [30, 20, 12, 6])
np.testing.assert_array_equal(daily_totals, [23, 23, 22])
assert product_totals.sum() == daily_totals.sum() == sales.sum() == 68

Expected output

Product totals: {'Books': 30, 'Pens': 20, 'Folders': 12, 'Gifts': 6}
Daily totals: [23 23 22]
Highest-total product: Books

Typical Axis Problems and Fixes 🔝

Unexpected shape: write the input shape and cross out the reduced axes. Wrong calculation runs: check the meaning of broadcasting, especially for square arrays. Constant feature or zero total: choose an explicit division policy. NaN-aware total looks like zero: disclose valid counts. Extreme location is ambiguous: argmax and argmin return the first position. Integer overflow: choose a bounded accumulator dtype. Axis out of range: inspect ndim; a one-dimensional array has no axis=1. Continue with shape-error troubleshooting and data types.

Three Exercises: Predict Shapes First 🔝

1. For a (6, 4) matrix, subtract row medians while retaining dimensions. 2. For a (3, 5) matrix, retrieve the maximum position and value per column. 3. For a (2, 3, 4) array, compute cumulative sums along the last axis. Predict output shapes (6, 4), (5,) for both positions and values, and (2, 3, 4).

Solution 1: Subtract Row Medians 🔝

The median result has shape (6, 1), so it broadcasts across each row. The row medians of the centered array should be zero.

M = np.arange(24).reshape(6, 4)
med = np.median(M, axis=1, keepdims=True)
centered = M - med
print(centered)
assert med.shape == (6, 1) and centered.shape == (6, 4)
np.testing.assert_allclose(np.median(centered, axis=1), 0)

Expected output

[[-1.5 -0.5  0.5  1.5]
 [-1.5 -0.5  0.5  1.5]
 [-1.5 -0.5  0.5  1.5]
 [-1.5 -0.5  0.5  1.5]
 [-1.5 -0.5  0.5  1.5]
 [-1.5 -0.5  0.5  1.5]]

Solution 2: Maximum Position per Column 🔝

Use a fixed array so the results are repeatable. Column IDs pair with the maximum row positions when selecting values. A tie returns the first row position.

A = np.array([[10, 50, 30, 90, 5], [20, 40, 80, 90, 10], [15, 60, 70, 20, 25]])
idx = A.argmax(axis=0)
vals = A[idx, np.arange(A.shape[1])]
print("Positions:", idx)
print("Values:", vals)
np.testing.assert_array_equal(idx, [1, 2, 1, 0, 2])
np.testing.assert_array_equal(vals, A.max(axis=0))

Expected output

Positions: [1 2 1 0 2]
Values: [20 60 80 90 25]

Solution 3: Cumulative Last-Axis Totals 🔝

Cumulative sums retain the original shape. The last value along each final-axis slice should equal the ordinary sum along that slice.

T = np.arange(24).reshape(2, 3, 4)
running = np.cumsum(T, axis=2)
print(running)
assert running.shape == T.shape
np.testing.assert_array_equal(running[:, :, -1], T.sum(axis=2))

Expected output

[[[ 0  1  3  6]
  [ 4  9 15 22]
  [ 8 17 27 38]]

 [[12 25 39 54]
  [16 33 51 70]
  [20 41 63 86]]]

Practice in Google Colab 🔝

Open in Google Colab View on GitHub
All synthetic sample data is included. Run the examples, predict reduced shapes and complete the three exercises. Save a copy in Drive before editing.

Continue with NumPy tutorials, sum(), mean(), std(), min(), max() and broadcasting. References: NumPy sum and NumPy mean documentation.

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: