sorting a tensor with the indices from argsort - is there a symbol operator?

[mureva]

Say I have some tensors:

• A of shape (b, s, 1)
• B of shape (b, s, n)

If I argsort along axis(1) I'll get suitable indices to also sort B. But... how can I use those indices? None of take, pick, or gather_nd seem to be suitable, so how can I actually use the sorting indices?

[mureva]

OK, so gather_nd can do it - but we need a fair bit of manipulation.
pseudocode:

i0 = argsort(A, axis=1)     // gives (bs,s,1)
i1 = transpose( i0, (2,0,1) // gives (1,bs,s)

ls0 = linspace(0,bs-1,bs)       // gives integer indices 0->(bs-1)
ls1 = reshape(ls0, (1,bs,1))
ls2 = broadcast_to(ls1, (1,bs,s))

inds = concat( {ls0,i1}, axis=0 )  // gives (2,bs,s)

sortedB = gather_nd( B, inds )