Reshape and Transpose - NumPy Fundamentals

Reshape and Transpose

What you'll build

A function reshape_transpose(arr) that reshapes a 24-element vector into 2D and 3D arrays, transposes them, and flattens back.

Background

Shape changes glue ML pipelines together: flattening images before a linear layer, splitting d_model into attention heads, moving channels between (N, H, W, C) and (N, C, H, W). Mixing up reshape (same order, new grouping) and transpose (same data, axes reordered) is one of the most common silent bugs.

Key ideas

• NumPy stores arrays in row-major (C) order: the last axis changes fastest.
• reshape keeps that order and regroups it. In a (4, 6) array, entry [i, j] is element 6*i + j of the vector; in a (2, 3, 4) array, entry [i, j, k] is element 12*i + 4*j + k. One dimension may be -1 ("infer it").
• transpose reorders axes: out = a.transpose(2, 1, 0) means out[k, j, i] = a[i, j, k]. For 2D, a.T[j, i] = a[i, j]. No data moves; only the strides change.
• flatten() reads an array back out in row-major order (always a copy); ravel() does the same but returns a view when it can.

Inputs

• arr: 1D array of exactly 24 elements (the tests mostly use np.arange(24)).

Expected Return Format

Return a dictionary with these keys:

• 'arr_2d': arr reshaped to (4, 6)
• 'arr_3d': arr reshaped to (2, 3, 4)
• 'arr_2d_transposed': transpose of arr_2d, shape (6, 4)
• 'arr_3d_swapped': arr_3d with its axes reversed, transpose(2, 1, 0), shape (4, 3, 2)
• 'arr_flat': arr_2d flattened back to 1D in row-major order, shape (24,) (equal to the input)

Edge cases and conventions

• arr_2d.reshape(6, 4) has the right shape but is not the transpose: its first row is [0, 1, 2, 3], while the transpose's first row is [0, 6, 12, 18].
• Flattening the transposed array gives a different order ([0, 6, 12, 18, 1, ...]), which is why 'arr_flat' flattens arr_2d.

Constraints

Use reshape, transpose/.T and flatten/ravel; no loops.

Follow-up

• Convert a batch of images from (N, H, W, C) to (N, C, H, W). Which axis order do you pass to transpose?
• Why does .view(...) on a transposed tensor fail in PyTorch, while .reshape(...) works?

Input: reshape_transpose(np.arange(24))["arr_2d_transposed"][:2]
Output: array([[ 0,  6, 12, 18],
       [ 1,  7, 13, 19]])

Input: reshape_transpose(np.arange(24))["arr_3d_swapped"][:, 0, :]
Output: array([[ 0, 12],
       [ 1, 13],
       [ 2, 14],
       [ 3, 15]])

Input: reshape_transpose(np.arange(24))["arr_3d"].shape
Output: (2, 3, 4)

import numpy as np

def reshape_transpose(arr: np.ndarray) -> dict:
    """
    Reshape, transpose and flatten a 24-element array.

    Args:
        arr: 1D array of 24 elements

    Returns:
        dict with keys:
            'arr_2d':            arr reshaped to (4, 6)
            'arr_3d':            arr reshaped to (2, 3, 4)
            'arr_2d_transposed': arr_2d transposed, shape (6, 4)
            'arr_3d_swapped':    arr_3d with axes reversed (2, 1, 0), shape (4, 3, 2)
            'arr_flat':          arr_2d flattened back to shape (24,)
    """
    # Your code here
    pass