(weights: NDArray, n_head: int, n_head_kv: int)
| 566 | |
| 567 | |
| 568 | def permute(weights: NDArray, n_head: int, n_head_kv: int) -> NDArray: |
| 569 | # print( "permute debug " + str(weights.shape[0]) + " x " + str(weights.shape[1]) + " nhead " + str(n_head) + " nheadkv " + str(n_kv_head) ) |
| 570 | if n_head_kv is not None and n_head != n_head_kv: |
| 571 | n_head = n_head_kv |
| 572 | return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:]) |
| 573 | .swapaxes(1, 2) |
| 574 | .reshape(weights.shape)) |
| 575 | |
| 576 | |
| 577 | class Tensor(metaclass=ABCMeta): |
no outgoing calls
no test coverage detected