kd.nn.PatchifyEmbed

kd.nn.PatchifyEmbed#

class kauldron.modules.PatchifyEmbed(patch_size: tuple[int, int], hidden_size: int, parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)[source]

Bases: flax.linen.module.Module

Patchify and linearly embed and image, as in ViT.

patch_size: tuple[int, int]
hidden_size: int
name: str | None = None
parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = None
scope: Scope | None = None