| # 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 | |