(self, E_q: int, E_k: int, E_v: int, E_total: int,
nheads: int, dropout_p: float = 0.0)
| 185 | dropout_p (float, optional): Dropout probability. Default: 0.0 |
| 186 | """ |
| 187 | def __init__(self, E_q: int, E_k: int, E_v: int, E_total: int, |
| 188 | nheads: int, dropout_p: float = 0.0): |
| 189 | super().__init__() |
| 190 | self.nheads = nheads |
| 191 | self.dropout_p = dropout_p |
| 192 | self.query_proj = nn.Linear(E_q, E_total) |
| 193 | self.key_proj = nn.Linear(E_k, E_total) |
| 194 | self.value_proj = nn.Linear(E_v, E_total) |
| 195 | E_out = E_q |
| 196 | self.out_proj = nn.Linear(E_total, E_out) |
| 197 | assert E_total % nheads == 0, "Embedding dim is not divisible by nheads" |
| 198 | self.E_head = E_total // nheads |
| 199 | |
| 200 | def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: |
| 201 | """ |
nothing calls this directly
no outgoing calls
no test coverage detected