Directional Aggregations with the axis Parameter0%

Sorting Arrays & Extracting Indices (np.sort vs np.argsort)

Directional Aggregations with the axis Parameter

Beginner12 min read•Updated: Jul 11, 2026

Mastering the axis Parameter: Reductions & keepdims

One of the most essential concepts in multi-dimensional computing is understanding how aggregation functions operate across specific axes.

Whether you are computing the average grade per student, the total sales per store across quarters, or normalizing batch activations in a convolutional neural network, mastering the axis and keepdims parameters is non-negotiable.

1. The Mental Model of the `axis` Parameter

A common misconception is thinking of axis as "operating along a dimension". A far clearer and more intuitive mental model is:

The axis parameter specifies the dimension that is COLLAPSED (eliminated) during the aggregation.

Consider a 2D matrix of shape (rows, cols) = (4, 3):

axis=0 (Rows collapsed): Collapses row-by-row down the columns. The row dimension disappears. Output shape is (3,) (one result per column).
axis=1 (Columns collapsed): Collapses column-by-column across the rows. The column dimension disappears. Output shape is (4,) (one result per row).
axis=None (Default): Collapses all dimensions, flattening the array into a single scalar value.

1import numpy as np
2 
3# 4 students (rows), 3 exam scores each (columns)
4grades = np.array([
5    [85, 90, 78],  # Student 0
6    [72, 88, 91],  # Student 1
7    [95, 92, 99],  # Student 2
8    [60, 70, 65]   # Student 3
9])
10 
11print("Original shape:", grades.shape)  # (4, 3)
12 
13# 1. Overall class average across all exams and all students (axis=None)
14print("Overall average:", np.mean(grades))  # 82.08
15 
16# 2. Average score for each exam subject (collapse rows -> axis=0)
17exam_averages = np.mean(grades, axis=0)
18print("Exam subject averages (shape:", exam_averages.shape, "):", exam_averages)
19# Output: [78.  85.  83.25] -> shape (3,)
20 
21# 3. Average score for each individual student (collapse cols -> axis=1)
22student_averages = np.mean(grades, axis=1)
23print("Student averages (shape:", student_averages.shape, "):", student_averages)
24# Output: [84.33  83.67  95.33  65.  ] -> shape (4,)

2. In 3D and N-Dimensional Tensors

The exact same rule holds for higher dimensions. If you have an array of shape (Batch, Height, Width, Channels) = (32, 28, 28, 3):

np.mean(tensor, axis=0) collapses the batch axis $\implies$ shape becomes (28, 28, 3).
np.mean(tensor, axis=(1, 2)) collapses spatial height and width $\implies$ shape becomes (32, 3).
np.mean(tensor, axis=-1) collapses the color channels $\implies$ shape becomes (32, 28, 28).

3. The Power of `keepdims=True` for Broadcasting

When you compute an axis reduction, NumPy collapses that dimension, dropping the rank by 1. For example, in a $(4, 3)$ matrix, grades.mean(axis=1) produces shape (4,) rather than (4, 1).

If you immediately try to subtract this mean from the original matrix to center the data: grades - student_averages NumPy attempts to align trailing dimensions: $(4, 3)$ with $(4,)$, which fails with a broadcasting ValueError!

Setting keepdims=True preserves the collapsed dimension as a singleton axis of length 1:

1# Calculate student average with keepdims=True -> shape becomes (4, 1)
2student_mean_2d = np.mean(grades, axis=1, keepdims=True)
3print("With keepdims=True shape:", student_mean_2d.shape)  # (4, 1)
4 
5# Now broadcasting works seamlessly! (4, 3) - (4, 1) -> (4, 3)
6centered_grades = grades - student_mean_2d
7print("
8Centered student grades (deviations from each student's average):
9", np.round(centered_grades, 2))

Output:

1With keepdims=True shape: (4, 1)
2 
3Centered student grades (deviations from each student's average):
4 [[ 0.67  5.67 -6.33]
5 [-11.67  4.33  7.33]
6 [-0.33 -3.33  3.67]
7 [-5.    5.    0.  ]]

4. Tuple of Axes Reductions

You can collapse multiple axes simultaneously by passing a tuple of integers:

1cube = np.ones((2, 5, 8))
2 
3# Collapse axis 0 and axis 2 simultaneously
4reduced = np.sum(cube, axis=(0, 2))
5print("Reduced shape:", reduced.shape)  # (5,)
6print("Values:", reduced)               # Each element is 2 * 8 = 16

Interactive Quiz• Knowledge Check

Multiple Choice Questions

Question 1

What is the output shape of np.sum(arr, axis=0) for an array of shape (10, 20)?

Question 2

In a 2D matrix of shape (N, M), which axis parameter computes row-wise sums (the sum across columns for each row)?

Question 3

What does setting keepdims=True do during an aggregation operation?

Question 4

If x has shape (4, 6) and you execute x - np.mean(x, axis=1), why does it raise a broadcasting ValueError?

Question 5

What is the output shape when applying np.mean(arr, axis=(0, 2)) to an array arr of shape (5, 8, 12)?

Next Lesson

Sorting Arrays & Extracting Indices (np.sort vs np.argsort)

Continue learning with hands-on practice, examples, and exercises in the upcoming topic.

Related Lessons

Previous Lesson	Next Lesson
Aggregate Statistics (sum, mean, median, std, var)	Sorting Arrays & Extracting Indices (np.sort vs np.argsort)

Practice Quiz

Test your understanding of this lesson with 5 questions. Each question has one correct answer.

Previous Next

PreviousPrevNext

MSK Institute • Shikohabad

Loading Page Content...

Preparing courses, live labs, and learning resources.सामग्री लोड हो रही है, कृपया प्रतीक्षा करें...

Connecting serverPlease wait...

💻 Practical Labs🚀 Free Demo📜 Verified Certificates

Sorting Arrays & Extracting Indices (np.sort vs np.argsort)

Directional Aggregations with the axis Parameter

Beginner12 min read•Updated: Jul 11, 2026

Study Materials Cheatsheet

Mastering the axis Parameter: Reductions & keepdims

One of the most essential concepts in multi-dimensional computing is understanding how aggregation functions operate across specific axes.

1. The Mental Model of the `axis` Parameter

A common misconception is thinking of axis as "operating along a dimension". A far clearer and more intuitive mental model is:

The axis parameter specifies the dimension that is COLLAPSED (eliminated) during the aggregation.

Consider a 2D matrix of shape (rows, cols) = (4, 3):

axis=0 (Rows collapsed): Collapses row-by-row down the columns. The row dimension disappears. Output shape is (3,) (one result per column).
axis=1 (Columns collapsed): Collapses column-by-column across the rows. The column dimension disappears. Output shape is (4,) (one result per row).
axis=None (Default): Collapses all dimensions, flattening the array into a single scalar value.

1import numpy as np
2 
3# 4 students (rows), 3 exam scores each (columns)
4grades = np.array([
5    [85, 90, 78],  # Student 0
6    [72, 88, 91],  # Student 1
7    [95, 92, 99],  # Student 2
8    [60, 70, 65]   # Student 3
9])
10 
11print("Original shape:", grades.shape)  # (4, 3)
12 
13# 1. Overall class average across all exams and all students (axis=None)
14print("Overall average:", np.mean(grades))  # 82.08
15 
16# 2. Average score for each exam subject (collapse rows -> axis=0)
17exam_averages = np.mean(grades, axis=0)
18print("Exam subject averages (shape:", exam_averages.shape, "):", exam_averages)
19# Output: [78.  85.  83.25] -> shape (3,)
20 
21# 3. Average score for each individual student (collapse cols -> axis=1)
22student_averages = np.mean(grades, axis=1)
23print("Student averages (shape:", student_averages.shape, "):", student_averages)
24# Output: [84.33  83.67  95.33  65.  ] -> shape (4,)

2. In 3D and N-Dimensional Tensors

The exact same rule holds for higher dimensions. If you have an array of shape (Batch, Height, Width, Channels) = (32, 28, 28, 3):

np.mean(tensor, axis=0) collapses the batch axis $\implies$ shape becomes (28, 28, 3).
np.mean(tensor, axis=(1, 2)) collapses spatial height and width $\implies$ shape becomes (32, 3).
np.mean(tensor, axis=-1) collapses the color channels $\implies$ shape becomes (32, 28, 28).

3. The Power of `keepdims=True` for Broadcasting

When you compute an axis reduction, NumPy collapses that dimension, dropping the rank by 1. For example, in a $(4, 3)$ matrix, grades.mean(axis=1) produces shape (4,) rather than (4, 1).

Setting keepdims=True preserves the collapsed dimension as a singleton axis of length 1:

1# Calculate student average with keepdims=True -> shape becomes (4, 1)
2student_mean_2d = np.mean(grades, axis=1, keepdims=True)
3print("With keepdims=True shape:", student_mean_2d.shape)  # (4, 1)
4 
5# Now broadcasting works seamlessly! (4, 3) - (4, 1) -> (4, 3)
6centered_grades = grades - student_mean_2d
7print("
8Centered student grades (deviations from each student's average):
9", np.round(centered_grades, 2))

Output:

1With keepdims=True shape: (4, 1)
2 
3Centered student grades (deviations from each student's average):
4 [[ 0.67  5.67 -6.33]
5 [-11.67  4.33  7.33]
6 [-0.33 -3.33  3.67]
7 [-5.    5.    0.  ]]

4. Tuple of Axes Reductions

You can collapse multiple axes simultaneously by passing a tuple of integers:

1cube = np.ones((2, 5, 8))
2 
3# Collapse axis 0 and axis 2 simultaneously
4reduced = np.sum(cube, axis=(0, 2))
5print("Reduced shape:", reduced.shape)  # (5,)
6print("Values:", reduced)               # Each element is 2 * 8 = 16

Interactive Quiz• Knowledge Check

Multiple Choice Questions

Question 1

What is the output shape of np.sum(arr, axis=0) for an array of shape (10, 20)?

Question 2

In a 2D matrix of shape (N, M), which axis parameter computes row-wise sums (the sum across columns for each row)?

Question 3

What does setting keepdims=True do during an aggregation operation?

Question 4

If x has shape (4, 6) and you execute x - np.mean(x, axis=1), why does it raise a broadcasting ValueError?

Question 5

What is the output shape when applying np.mean(arr, axis=(0, 2)) to an array arr of shape (5, 8, 12)?

Next Lesson

Sorting Arrays & Extracting Indices (np.sort vs np.argsort)

Continue learning with hands-on practice, examples, and exercises in the upcoming topic.

Related Lessons

Previous Lesson	Next Lesson
Aggregate Statistics (sum, mean, median, std, var)	Sorting Arrays & Extracting Indices (np.sort vs np.argsort)

Practice Quiz

Test your understanding of this lesson with 5 questions. Each question has one correct answer.

Previous Next

Mastering the axis Parameter: Reductions & keepdims

1. The Mental Model of the axis Parameter

2. In 3D and N-Dimensional Tensors

3. The Power of keepdims=True for Broadcasting

4. Tuple of Axes Reductions

Multiple Choice Questions

Sorting Arrays & Extracting Indices (np.sort vs np.argsort)

Related Lessons

Practice Quiz

Loading Page Content...

Mastering the axis Parameter: Reductions & keepdims

1. The Mental Model of the axis Parameter

2. In 3D and N-Dimensional Tensors

3. The Power of keepdims=True for Broadcasting

4. Tuple of Axes Reductions

Multiple Choice Questions

Sorting Arrays & Extracting Indices (np.sort vs np.argsort)

Related Lessons

Practice Quiz

1. The Mental Model of the `axis` Parameter

3. The Power of `keepdims=True` for Broadcasting

1. The Mental Model of the `axis` Parameter

3. The Power of `keepdims=True` for Broadcasting