-
Notifications
You must be signed in to change notification settings - Fork 27
Closed
Description
Line 206 in acd78b4
query_states, key_states = apply_rotary_pos_emb( |
the apply_rotary_pos_emb() function does not accept the offset argument?
def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
# The first two dimensions of cos and sin are always 1, so we can squeeze
them.
cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
q_embed = (q * cos) + (rotate_half(q) * sin)
k_embed = (k * cos) + (rotate_half(k) * sin)
return q_embed, k_embed
Metadata
Metadata
Assignees
Labels
No labels