Is there a good alternative to the method .apply_()?

1

Hi everyone,
At some point in one of my codes, I’d like to transform a tensor T of size (N x N x M x M) to a tensor U of size (A x N x N x M) such as U_ajkl = T_jklf(a, i) where f(a, i) is a function whose values are stored in a tensor F for the moment.

So with python loops it simply gives:

for a in range(A):
     for j in range(N):
         for k in range(N):
             for l in range(M):
                 U[a, j, k, l] = T[j, k, l, F[a, i]]

Obviously, I’d like to avoid such loops. I found the method .apply_() that may do the job but the doc says that it works only on CPU and that this function should not be used in sections that require high perfomance (which is my case). Have you any idea how to do better?