Dot product in EZKL

[biancaganescu]

Getting ezkl error when using dot product in self-attention layer forward implemented as follows:

def forward(self, x):
      B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

      # calculate query, key, values for all heads in batch and move head forward to be the batch dim
      q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
      k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
      q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
      v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

      # manual implementation of attention
      att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) #q shape:(B, nh, T, hs), k transpose shape (B, nh, hs, T) -> (B, nh, T, T)
      att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
      att = F.softmax(att, dim=-1)
      att = self.attn_dropout(att)
      y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
      y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

      y = self.resid_dropout(self.c_proj(y))
      return y

where c_attn and c_proj are linear layers.

[alexander-camuto]

Resolved by latest einsum updates :)