(x, indices)
| 956 | return w_hat, qw, s, b |
| 957 | |
| 958 | def gather_sort(x, indices): |
| 959 | N, M = indices.shape |
| 960 | indices = indices.flatten() |
| 961 | order = mx.argsort(indices) |
| 962 | inv_order = mx.argsort(order) |
| 963 | return x.flatten(0, -3)[order // M], indices[order], inv_order |
| 964 | |
| 965 | def scatter_unsort(x, inv_order, shape=None): |
| 966 | x = x[inv_order] |
nothing calls this directly
no outgoing calls
no test coverage detected