In Tensorflow is there a way to permute batched tensors using scatter_nd without unstacking the tensor?

Question

For example, lets say my input tensor is of dimensionss (bs, nv, dims).

I have another tensor of indices of shape (bs, nv)

What I wou