PatchEmbedding

A layer that convert image into patch embeddings.

PatchEmbedding(patch_size=16, channels=3, dim=768)

Forward

(x: jaxtyping.Float[Tensor, '... c h w']) -> jaxtyping.Float[Tensor, '... n d']