| 375 | return tensor.reshape(bsz, seq_len, self.num_heads, self.head_dim).permute(0,2,1,3) |
| 376 | |
| 377 | def __call__(self, hidden_states, causal_attention_mask): |
| 378 | bsz, tgt_len, embed_dim = hidden_states.shape |
| 379 | |
| 380 | query_states = self.q_proj(hidden_states) * self.scale |
| 381 | key_states = self._shape(self.k_proj(hidden_states), -1, bsz) |
| 382 | value_states = self._shape(self.v_proj(hidden_states), -1, bsz) |
| 383 | |
| 384 | proj_shape = (bsz * self.num_heads, -1, self.head_dim) |
| 385 | query_states = self._shape(query_states, tgt_len, bsz).reshape(*proj_shape) |
| 386 | key_states = key_states.reshape(*proj_shape) |
| 387 | src_len = key_states.shape[1] |
| 388 | value_states = value_states.reshape(*proj_shape) |
| 389 | |
| 390 | attn_weights = query_states @ key_states.permute(0,2,1) |
| 391 | |
| 392 | attn_weights = attn_weights.reshape(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask |
| 393 | attn_weights = attn_weights.reshape(bsz * self.num_heads, tgt_len, src_len) |
| 394 | |
| 395 | attn_weights = attn_weights.softmax() |
| 396 | |
| 397 | attn_output = attn_weights @ value_states |
| 398 | |
| 399 | attn_output = attn_output.reshape(bsz, self.num_heads, tgt_len, self.head_dim) |
| 400 | attn_output = attn_output.permute(0,2,1,3) |
| 401 | attn_output = attn_output.reshape(bsz, tgt_len, embed_dim) |
| 402 | |
| 403 | attn_output = self.out_proj(attn_output) |
| 404 | return attn_output |
| 405 | |
| 406 | class CLIPEncoderLayer: |
| 407 | def __init__(self): |