import numpy as np
from numpy.typing import NDArray


def patchify(image: NDArray[np.float64], patch_size: int) -> NDArray[np.float64]:
    h, w, c = image.shape
    p = patch_size
    if p <= 0 or h % p or w % p:
        raise ValueError("Patch size must divide height and width")
    grid = image.reshape(h // p, p, w // p, p, c)
    patches = grid.transpose(0, 2, 1, 3, 4)
    return patches.reshape((h // p) * (w // p), p * p * c)


image = np.arange(224 * 224 * 3, dtype=np.float64).reshape(224, 224, 3)
patches = patchify(image, patch_size=16)

assert patches.shape == (196, 768)
np.testing.assert_array_equal(
    patches[0], image[:16, :16, :].reshape(-1)
)
np.testing.assert_array_equal(
    patches[1], image[:16, 16:32, :].reshape(-1)
)
print(patches.shape)  # (196, 768)