Transpose

Transpose moves values according to axis meaning. You tell it the new order of the axes, and it rearranges accordingly.

chw   = np.transpose(photo, (2, 0, 1))      # (3, 512, 512) — correct
batch = np.stack([photo, photo, photo])      # (3, 512, 512, 3)
nchw  = np.transpose(batch, (0, 3, 1, 2))    # (3, 3, 512, 512)

The tuple is read as “the new axis 0 is the old axis 2, the new axis 1 is the old axis 0…”.