| 657 | |
| 658 | |
| 659 | def _get_idx(index, axis): |
| 660 | index_dims = len(index.shape) |
| 661 | idx = [] |
| 662 | if axis < 0: |
| 663 | axis += index_dims |
| 664 | for i in range(index_dims): |
| 665 | if i != axis: |
| 666 | shape = [1] * index_dims |
| 667 | shape[i] = index.shape[i] |
| 668 | arange = linspace( |
| 669 | 0, index.shape[i] - 1, index.shape[i], device=index.device, |
| 670 | ) |
| 671 | arange = ( |
| 672 | broadcast_to(arange.reshape(*shape), index.shape) |
| 673 | .reshape(-1) |
| 674 | .astype(np.int32) |
| 675 | ) |
| 676 | idx.append(arange) |
| 677 | else: |
| 678 | idx.append(index.reshape(-1)) |
| 679 | return tuple(idx) |
| 680 | |
| 681 | |
| 682 | def gather(inp: Tensor, axis: int, index: Tensor) -> Tensor: |