(lazy_tensor: LazyTensor, n_head: int, n_head_kv: int)
| 736 | |
| 737 | |
| 738 | def permute_lazy(lazy_tensor: LazyTensor, n_head: int, n_head_kv: int) -> LazyTensor: |
| 739 | def load() -> Tensor: |
| 740 | return lazy_tensor.load().permute(n_head, n_head_kv) |
| 741 | return LazyTensor(load, lazy_tensor.shape, lazy_tensor.data_type, f'permute({n_head}, {n_head_kv}) ' + lazy_tensor.description) |
| 742 | |
| 743 | |
| 744 | def permute_part_lazy(lazy_tensor: LazyTensor, n_part: int, n_head: int, n_head_kv: int) -> LazyTensor: |
no test coverage detected