(weights: NDArray, n_head: int, n_head_kv: int)
| 692 | |
| 693 | |
| 694 | def permute(weights: NDArray, n_head: int, n_head_kv: int) -> NDArray: |
| 695 | if n_head_kv is not None and n_head != n_head_kv: |
| 696 | n_head = n_head_kv |
| 697 | return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:]) |
| 698 | .swapaxes(1, 2) |
| 699 | .reshape(weights.shape)) |
| 700 | |
| 701 | |
| 702 | class Tensor(ABC): |
no outgoing calls
no test coverage detected