File size: 596 Bytes
629b123 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | # SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
import torch
def stack_cos_sin(cos, sin):
cos = torch.stack([cos, cos], dim=-1).flatten(-2)
sin = torch.stack([sin, sin], dim=-1).flatten(-2)
return cos, sin
def get_rot_transformation_mat():
# ROPE op uses a single tile
dhead = 32
rot_emb_matrix = torch.zeros(1, 1, dhead, dhead)
rot_emb_matrix[..., torch.arange(0, dhead, 2), torch.arange(1, dhead, 2)] = 1
rot_emb_matrix[..., torch.arange(1, dhead, 2), torch.arange(0, dhead, 2)] = -1
return rot_emb_matrix
|