How do I swap tensor's axes in TensorFlow?

tf.transpose provides the same functionality as np.swapaxes, although in a more generalized form. In your case, you can do tf.transpose(orig_tensor, [1, 0, 2]) which would be equivalent to np.swapaxes(orig_np_array, 0, 1).

It is possible to use tf.einsum to swap axes if the number of input dimensions is unknown. For example:

• tf.einsum("ij...->ji...", input) will swap the first two dimensions of input;
• tf.einsum("...ij->...ji", input) will swap the last two dimensions;
• tf.einsum("aij...->aji...", input) will swap the second and the third dimension;
• tf.einsum("ijk...->kij...", input) will permute the first three dimensions;

and so on.