Applying a permutation along one axis in TensorFlow
tf.gather can be used to that end. In fact, it is even more general, as the indices it takes as one of its inputs don't need to represent a permutation.
How to permute "dimensions" along a single axis of a tensor?
Something akin to tf.transpose, but at the level of "dimensions" along an axis, instead of at the level of axes.
To permute them randomly (along the first axis), there it tf.random.shuffle, and to shift them, there is tf.roll. But I can't find a more general function that would apply any given permutation.