numpy
multidimensional arrays
array manipulation
python
flattening arrays

How to flatten only some dimensions of a numpy array

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Flattening only some dimensions of a NumPy array really means reshaping selected axes into one combined axis while leaving the others intact. The operation is straightforward once you think in terms of shape algebra instead of “flatten” as an all-or-nothing action.

Flatten Contiguous Axes With reshape

If the axes you want to combine are adjacent, reshape is usually enough. For example, suppose an array has shape (2, 3, 4) and you want to merge the last two axes into one axis of size 12.

python
1import numpy as np
2
3arr = np.arange(24).reshape(2, 3, 4)
4result = arr.reshape(2, 12)
5
6print(arr.shape)
7print(result.shape)
8print(result)

The logic is simple: 3 * 4 = 12, so the new shape becomes (2, 12).

You can do the same with the first two axes:

python
result = arr.reshape(6, 4)
print(result.shape)

Here 2 * 3 = 6, so the result shape is (6, 4).

Use -1 When One Dimension Is Implied

NumPy can infer one dimension automatically with -1, which makes shape changes less error-prone.

python
result = arr.reshape(arr.shape[0], -1)
print(result.shape)

This says “keep the first axis as it is and flatten everything else.” It is especially useful when the exact shape is not hardcoded.

Non-Adjacent Axes Need Reordering First

If the axes you want to combine are not next to each other, move them together first with transpose or moveaxis, then reshape.

Suppose the shape is (2, 3, 4, 5) and you want to combine axes 0 and 2 while leaving the others logically separate. Since those axes are not adjacent, reorder them first.

python
1arr = np.arange(120).reshape(2, 3, 4, 5)
2reordered = np.transpose(arr, (0, 2, 1, 3))
3result = reordered.reshape(8, 3, 5)
4
5print(arr.shape)
6print(reordered.shape)
7print(result.shape)

After reordering, axes 0 and 2 from the original array sit next to each other as the first two axes in reordered, so they can be merged cleanly into 8.

A Practical Mental Model

A useful way to reason about the operation is:

  1. decide the axis order you want
  2. move those axes into adjacent positions if necessary
  3. multiply the sizes of the axes you want to combine
  4. call reshape

Once you do that, selective flattening is just controlled reshaping.

Views Versus Copies

reshape often returns a view when memory layout permits it, which is efficient. But after certain transposes or slices, the array may no longer be contiguous in the way reshape expects, and NumPy may need to create a copy.

You can inspect this with the flags attribute.

python
1arr = np.arange(24).reshape(2, 3, 4)
2reordered = arr.transpose(1, 0, 2)
3reshaped = reordered.reshape(3, 8)
4
5print(reordered.flags['C_CONTIGUOUS'])
6print(reshaped.flags['C_CONTIGUOUS'])

For correctness, this usually does not matter. For large arrays, it can matter a lot for memory use and speed.

A Reusable Helper Function

If you perform this kind of operation often, a helper can make the intent clearer.

python
1import numpy as np
2from math import prod
3
4
5def flatten_axes(arr, start_axis, end_axis):
6    shape = arr.shape
7    merged = prod(shape[start_axis:end_axis + 1])
8    new_shape = shape[:start_axis] + (merged,) + shape[end_axis + 1:]
9    return arr.reshape(new_shape)
10
11
12arr = np.arange(120).reshape(2, 3, 4, 5)
13result = flatten_axes(arr, 1, 2)
14print(result.shape)

This function flattens a contiguous block of axes, which covers many practical tensor reshaping tasks.

Common Use Cases

Selective flattening appears often in scientific and machine learning code:

  • combine spatial dimensions before feeding data into a dense layer
  • collapse batch and time axes for vectorized processing
  • flatten image height and width while preserving channel count
  • reshape grouped measurements into a 2D analysis matrix

The operation is not special to NumPy. It is a general tensor manipulation pattern.

Common Pitfalls

A common mistake is using flatten() when you only meant to merge some axes. flatten() always returns a fully 1D copy.

Another mistake is forgetting that non-adjacent axes cannot be merged directly without reordering the shape first.

Developers also miscompute the target size manually. Let NumPy infer one dimension with -1 where possible.

Finally, watch memory layout when working with very large arrays. A transpose followed by reshape may create a copy, which can be expensive.

Summary

  • Flattening some dimensions is usually just reshape on selected axes.
  • Merge adjacent axes directly by multiplying their sizes.
  • Reorder non-adjacent axes first with transpose or moveaxis.
  • Use -1 when NumPy can infer one dimension safely.
  • Be aware that some reshape operations after transposition may allocate a copy.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.