(lazy_tensor: LazyTensor, n_head: int, n_head_kv: int)
| 922 | |
| 923 | |
| 924 | def permute_lazy(lazy_tensor: LazyTensor, n_head: int, n_head_kv: int) -> LazyTensor: |
| 925 | def load() -> Tensor: |
| 926 | return lazy_tensor.load().permute(n_head, n_head_kv) |
| 927 | return LazyTensor(load, lazy_tensor.shape, lazy_tensor.data_type, f'permute({n_head}, {n_head_kv}) ' + lazy_tensor.description) |
| 928 | |
| 929 | |
| 930 | def permute_part_lazy(lazy_tensor: LazyTensor, n_part: int, n_head: int, n_head_kv: int) -> LazyTensor: |
no test coverage detected