| 607 | return x |
| 608 | |
| 609 | class TokenSplitWithoutSkip(nn.Module): |
| 610 | def __init__(self, in_features, out_features, patch_size=(2, 2)): |
| 611 | super().__init__() |
| 612 | self.h = patch_size[0] |
| 613 | self.w = patch_size[1] |
| 614 | self.proj = apply_wd(Linear(in_features, out_features * self.h * self.w, bias=False)) |
| 615 | nn.init.zeros_(self.proj.weight) |
| 616 | |
| 617 | def forward(self, x): |
| 618 | x = self.proj(x) |
| 619 | return rearrange(x, "... h w (nh nw e) -> ... (h nh) (w nw) e", nh=self.h, nw=self.w) |
| 620 | |
| 621 | class TokenSplit(nn.Module): |
| 622 | def __init__(self, in_features, out_features, patch_size=(2, 2)): |