How can we represent images, like sentences, words, and letters?

def patchify(x: torch.Tensor, patch_size: int) -> torch.Tensor:
    """Convert images to flattened non-overlapping patch tokens.

    Args:
        x: Image batch with shape (batch_size, channels, height, width).
        patch_size: Side length of each square patch.

    Returns:
        Tensor with shape (batch_size, num_patches, channels * patch_size * patch_size).
    """
    batch_size, channels, height, width = x.shape

    if height % patch_size != 0 or width % patch_size != 0:
        raise ValueError("Image height and width must be divisible by patch_size")

    patches_h = height // patch_size
    patches_w = width // patch_size

    x = x.reshape(batch_size, channels, patches_h, patch_size, patches_w, patch_size)
    x = x.permute(0, 2, 4, 1, 3, 5).contiguous()  # permute() changes view order; contiguous() fixes memory layout
    return x.reshape(batch_size, patches_h * patches_w, channels * patch_size * patch_size)

# Patch projection and learned positional encoding from ViT.
class PatchEmbed(nn.Module):
    def __init__(self, patch_dim: int, d_model: int): #linear projection 
        super().__init__()
        self.proj = nn.Linear(patch_dim, d_model)

    def forward(self, x_patches: torch.Tensor) -> torch.Tensor:
        """Project flattened patches from patch_dim to d_model."""
        return self.proj(x_patches)


class PositionalEmbedding(nn.Module):
    def __init__(self, num_tokens: int, d_model: int):
        super().__init__()
        self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_tokens + 1, d_model))

        #fill this tensor with values from a normal distribution, but truncate extreme values.
        #std=0.02 means the random values are very small, close to zero
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        nn.init.trunc_normal_(self.pos_embed, std=0.02)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Prepend the class token and add learned positional embeddings."""
        batch_size = x.shape[0]
        cls_tokens = self.cls_token.expand(batch_size, -1, -1) #-1 means same dimension
        #concatenate with the class token x.shape = (batch_size, num_patches + 1, d_model)
        x = torch.cat([cls_tokens, x], dim=1)
        return x + self.pos_embed[:, : x.shape[1], :] #position vectors are added to patches in the batch

# FFN baseline and GLU variants to compare in the Transformer blocks.
class FeedForward(nn.Module):
    """
    Standard Transformer FFN:
      x -> Linear(d_model->d_ff) -> GELU -> Dropout -> Linear(d_ff->d_model) -> Dropout
    """
    def __init__(self, d_model: int, d_ff: int, dropout: float):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout), #dropping out some percent of neurons during the training phase
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class GLUFeedForward(nn.Module):
    """GLU-family FFN: act(W_a x) * W_b x, followed by projection back to d_model."""
    def __init__(self, d_model: int, d_ff_gated: int, dropout: float, variant: str):
        super().__init__()
        self.variant = variant.lower()
        self.proj_in = nn.Linear(d_model, 2 * d_ff_gated) #self.proj_in(x).shape == (batch_size, tokens, 512)
        self.dropout = nn.Dropout(dropout)
        self.proj_out = nn.Linear(d_ff_gated, d_model)

        valid_variants = {"geglu", "swiglu", "reglu", "glu"}
        if self.variant not in valid_variants:
            raise ValueError(f"Unknown GLU variant '{variant}'. Choose one of {sorted(valid_variants)}.")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        a, gate = self.proj_in(x).chunk(2, dim=-1) #split the tensor into two parts. take the last batch to GLU
        #what GLU does activation(a) * gate

#activation functions
        if self.variant == "geglu":
            a = F.gelu(a)
        elif self.variant == "swiglu":
            a = F.silu(a)
        elif self.variant == "reglu":
            a = F.relu(a)
        elif self.variant == "glu":
            a = torch.sigmoid(a)
        #proj_in -> split -> activation/gate multiply -> Dropout -> proj_out -> Dropout
        x = a * gate
        x = self.dropout(x) #drop hidden features before output projection
        x = self.proj_out(x)
        return self.dropout(x) #final out put of the MLP before residual addition

class TransformerEncoderBlock(nn.Module):
    """
    Pre-LN encoder block: #LN means LayerNorm, normalizing across the feature dimension 
      x = x + Dropout(SelfAttn(LN(x)))
      x = x + Dropout(MLP(LN(x)))
    """
    #basically following the architecture of the paper 
    def __init__(self, d_model: int, n_heads: int, mlp: nn.Module, dropout: float):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention( #nn.MultiheadAttention expects query, key and value
            embed_dim=d_model,
            num_heads=n_heads,
            dropout=dropout,
            batch_first=True,
        )
        self.dropout = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.mlp = mlp

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        attn_input = self.norm1(x)
        attn_out, _ = self.attn(attn_input, attn_input, attn_input, need_weights=False) #for query, key and value
        x = x + self.dropout(attn_out)
        x = x + self.mlp(self.norm2(x))
        return x