(*, weight, bias, split, chunk_size)
| 676 | |
| 677 | # TODO maybe document and/or can do more efficiently (build indices in for loop and extract once for each split?) |
| 678 | def split_attentions(*, weight, bias, split, chunk_size): |
| 679 | weights = [None] * split |
| 680 | biases = [None] * split |
| 681 | |
| 682 | weights_biases_idx = 0 |
| 683 | |
| 684 | for starting_row_index in range(0, weight.shape[0], chunk_size): |
| 685 | row_indices = torch.arange(starting_row_index, starting_row_index + chunk_size) |
| 686 | |
| 687 | weight_rows = weight[row_indices, :] |
| 688 | bias_rows = bias[row_indices] |
| 689 | |
| 690 | if weights[weights_biases_idx] is None: |
| 691 | weights[weights_biases_idx] = weight_rows |
| 692 | biases[weights_biases_idx] = bias_rows |
| 693 | else: |
| 694 | assert weights[weights_biases_idx] is not None |
| 695 | weights[weights_biases_idx] = torch.concat([weights[weights_biases_idx], weight_rows]) |
| 696 | biases[weights_biases_idx] = torch.concat([biases[weights_biases_idx], bias_rows]) |
| 697 | |
| 698 | weights_biases_idx = (weights_biases_idx + 1) % split |
| 699 | |
| 700 | return weights, biases |
| 701 | |
| 702 | |
| 703 | def parse_list(value): |
no outgoing calls
no test coverage detected