Implement Vision Transformer (ViT) patch embedding, which turns an image into a sequence of patch vectors for transformer input. You’ll split the image into non-overlapping patches, flatten each patch, then apply a learned linear projection to get embeddings.

The patch embedding can be written as:

where (X \in \mathbb{R}^{N \times (P^2C)}) is the matrix of flattened patches, (W \in \mathbb{R}^{(P^2C) \times D}) is the projection matrix, (b \in \mathbb{R}^{D}), (N) is the number of patches, (P) is the patch size, (C) is the number of channels, and (D) is the embedding dimension.

Requirements

Implement the function

Rules:

• Split image into non-overlapping patches of size patch_size x patch_size.
• Flatten each patch in row-major order, keeping channel values last (i.e., flatten over H, then W, then C within a patch).
• Stack all flattened patches into a matrix X with shape (N, P*P*C) in top-to-bottom, left-to-right patch order.
• Compute embeddings using E = X @ W + b (use NumPy for matrix math).
• Return E as a NumPy array.

Example

Output: