| 703 | |
| 704 | |
| 705 | class PackedInputProjection(nn.Module): |
| 706 | def __init__(self, E_q, E_total, bias=False, device=None, dtype=None): |
| 707 | factory_kwargs = {"device": device, "dtype": dtype} |
| 708 | super().__init__() |
| 709 | self.packed_proj = nn.Linear(E_q, E_total * 3, bias=bias, **factory_kwargs) |
| 710 | |
| 711 | def forward(self, query): |
| 712 | return torch.chunk(self.packed_proj(query), 3, dim=-1) |
| 713 | |
| 714 | |
| 715 | B, D, dtype = 256, 8192, torch.bfloat16 |
no outgoing calls
no test coverage detected