- def forward(self, x):
- q = torch.einsum('ntc,hdc->nhtd', x, self.wq)
- k = torch.einsum('ntc,hdc->nhtd', x, self.wk)
- v = torch.einsum('ntc,hdc->nhtd', x, self.wv)
- r = math.sqrt(q.size(3))
- a = torch.einsum('nhtd,nhsd->nhts', q, k).div(r)
+ self.w_q = randw(nb_heads, dim_qk, dim_in)
+ self.w_k = randw(nb_heads, dim_qk, dim_in)
+ self.w_v = randw(nb_heads, dim_v, dim_in)
+ self.w_o = randw(dim_v * nb_heads, dim_in)
+
+ def forward(self, x_q, x_kv = None):
+ if x_kv is None: x_kv = x_q
+
+ q = torch.einsum('ntc,hdc->nhtd', x_q, self.w_q)
+ k = torch.einsum('ntc,hdc->nhtd', x_kv, self.w_k)
+ v = torch.einsum('ntc,hdc->nhtd', x_kv, self.w_v)
+
+ a = torch.einsum('nhtd,nhsd->nhts', q, k) / math.sqrt(q.size(3))
+