Sort `x` along `dim` and apply the same permutation to `ids`. Args: x: values tensor (e.g. logits) ids: index tensor, same shape as x (e.g. vocab offsets) dim: dimension to sort along (only last dim supported) descending: sort order Returns: (sorted_
(
x,
ids,
dim: core.constexpr = None,
descending: core.constexpr = core.CONSTEXPR_0,
)
| 13 | |
| 14 | @triton.jit |
| 15 | def argsort( |
| 16 | x, |
| 17 | ids, |
| 18 | dim: core.constexpr = None, |
| 19 | descending: core.constexpr = core.CONSTEXPR_0, |
| 20 | ): |
| 21 | """Sort `x` along `dim` and apply the same permutation to `ids`. |
| 22 | |
| 23 | Args: |
| 24 | x: values tensor (e.g. logits) |
| 25 | ids: index tensor, same shape as x (e.g. vocab offsets) |
| 26 | dim: dimension to sort along (only last dim supported) |
| 27 | descending: sort order |
| 28 | |
| 29 | Returns: |
| 30 | (sorted_x, permuted_ids) |
| 31 | """ |
| 32 | _dim: core.constexpr = len(x.shape) - 1 if dim is None else dim |
| 33 | core.static_assert( |
| 34 | _dim == len(x.shape) - 1, |
| 35 | "only minor dimension is currently supported", |
| 36 | ) |
| 37 | log_n: core.constexpr = _log2(x.shape[_dim]) |
| 38 | n_dims: core.constexpr = _log2(x.numel) |
| 39 | |
| 40 | # Reshape to hypercube of shape [2, 2, ..., 2] |
| 41 | hx = core.reshape(x, [2] * n_dims if n_dims else [1]) |
| 42 | hi = core.reshape(ids, [2] * n_dims if n_dims else [1]) |
| 43 | |
| 44 | # Run only log_n stages (sort-axis size), not n_dims (total). |
| 45 | # The alternating flip pattern creates independent sorting networks per row. |
| 46 | for stage in core.static_range(1, log_n + 1): |
| 47 | hx, hi = _bitonic_merge(hx, hi, stage, 2 if stage < log_n else descending, n_dims) |
| 48 | |
| 49 | x = core.reshape(hx, x.shape) |
| 50 | ids = core.reshape(hi, ids.shape) |
| 51 | return x, ids |
| 52 | |
| 53 | |
| 54 | @triton.jit |