(lazy_tensor: LazyTensor, n_part: int, n_head: int, n_head_kv: int)
| 928 | |
| 929 | |
| 930 | def permute_part_lazy(lazy_tensor: LazyTensor, n_part: int, n_head: int, n_head_kv: int) -> LazyTensor: |
| 931 | def load() -> Tensor: |
| 932 | return lazy_tensor.load().permute_part(n_part, n_head, n_head_kv) |
| 933 | s = lazy_tensor.shape.copy() |
| 934 | s[0] = s[0] // 3 |
| 935 | return LazyTensor(load, s, lazy_tensor.data_type, f'permute({n_head}, {n_head_kv}) ' + lazy_tensor.description) |
| 936 | |
| 937 | |
| 938 | def part_lazy(lazy_tensor: LazyTensor, n_part: int) -> LazyTensor: |
no test coverage detected