Transpose of a 3D tensor

Question

I need to transpose a 3-dimensional tensor of the shape (batch_size, N, M) to (batch_size, M, N) in a custom loss function in Keras with tensorflow as the backend. I tried using the following function

K.permute_dimensions(y_true,(0,2,1))

It gives me an error that the dimension must be 2 instead, 3 is given.

When I try to input the following

K.permute_dimensions(y_true,(1,0))

it says the dimension is 3, instead, 2 is given.

I also tried using K.transpose but it gives me (M,N,batch_size).

I already implemented the solution with for loop, but in keras custom loss, for loops should be avoided.

from keras import backend as K
def custom_loss(y_true, y_pred):
    beta  = K.permute_dimensions(y_pred, (0,2,1))

Answer 1

Use tf.transpose instead:

tf.transpose(y_true, perm=(0, 2, 1))

Otherwise use tf.keras instead of simply keras, as follows:

from tensorflow.keras import backend as K

def custom_loss(y_true, y_pred):
    beta = K.permute_dimensions(y_pred, (0, 2, 1))
    # ...