How do I swap tensor's axes in TensorFlow?

Question

I have a tensor of shape (30, 116, 10), and I want to swap the first two dimensions, so that I have a tensor of shape (116, 30, 10)

I saw that numpy as such a function implemented (np.swapaxes) and I searched for something similar in tensorflow but I found nothing.

Do you have any idea?

Solution 1

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).

Solution 2

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.

Solution 3

You can transpose the last two axes only with tf.linalg.matrix_transpose, or more generally, you can swap any number of trailing axes by working out what the leading indices are dynamically, and using relative indices for the axes you want to transpose

x = tf.ones([5, 3, 7, 11])
trailing_axes = [-1, -2]

leading = tf.range(tf.rank(x) - len(trailing_axes))   # [0, 1]
trailing = trailing_axes + tf.rank(x)                 # [3, 2]
new_order = tf.concat([leading, trailing], axis=0)    # [0, 1, 3, 2]
res = tf.transpose(x, new_order)
res.shape                                             # [5, 3, 11, 7]