| 691 | |
| 692 | |
| 693 | class InputProjection(nn.Module): |
| 694 | def __init__(self, E_q, E_total, bias=False, device=None, dtype=None): |
| 695 | factory_kwargs = {"device": device, "dtype": dtype} |
| 696 | super().__init__() |
| 697 | self.q_proj = nn.Linear(E_q, E_total, bias=bias, **factory_kwargs) |
| 698 | self.k_proj = nn.Linear(E_q, E_total, bias=bias, **factory_kwargs) |
| 699 | self.v_proj = nn.Linear(E_q, E_total, bias=bias, **factory_kwargs) |
| 700 | |
| 701 | def forward(self, x): |
| 702 | return self.q_proj(x), self.k_proj(x), self.v_proj(x) |
| 703 | |
| 704 | |
| 705 | class PackedInputProjection(nn.Module): |
no outgoing calls
no test coverage detected